Ошибка CUDA из памяти при попытке реализовать тройную потерю

Question

Я пытаюсь реализовать тройную потерю для слоя NetVLAD, используя 3 разных изображения из моего загрузчика данных следующим образом:

• а - партия с минимальной прибавкой
• p — та же партия с большим количеством дополнений
• n — другая партия из одного и того же загрузчика данных

Однако когда я пытаюсь получить внедрения VLAD из моей модели, я получаю ошибку нехватки памяти после вычисления первого внедрения VLAD. Есть ли более эффективный способ запустить этот код?

РЕДАКТИРОВАТЬ: Возможно, лучший способ загрузить другую партию, которая мне нужна для негативов?

      for i, (inputs, poses) in enumerate(train_loader):

    inputs = inputs.to(device)
    # poses = poses.to(device)

    print(inputs.shape)  # anchors
    positives = positive_tfs(inputs)
    positives = positives.to(device)
    print(positives.shape)
    negatives, _ = next(iter(train_loader))  # different random batch thanks to shuffle=True
    negatives = negatives.to(device)
    print(negatives.shape)

    # Zero the parameter gradient
    optimizer.zero_grad()

    # FORWARD
    a_vlad, pos_out, ori_out = model(inputs)
    print(a_vlad.shape)
    p_vlad = model(positives, get_pose=False)
    print(p_vlad.shape)
    n_vlad = model(negatives, get_pose=False)
    print(n_vlad.shape)

    # COMPUTE LOSS
    loss = triplet_loss(a_vlad, p_vlad, n_vlad)
    print(loss)

Answer 1

Лучшее, что вы можете сделать, это попытаться очистить кеш. Это часто освобождает место и решает проблему. Добавьте это в начало вашего кода.

      import torch
torch.cuda.empty_cache()

Если это не помогло, попробуйте уменьшить размер пакета или модели.

Если нет, то, к сожалению, вы ничего не сможете сделать, если на вашем графическом процессоре недостаточно памяти...

Вы можете посмотреть здесь: Как избежать «CUDA нехватки памяти» в PyTorch.