Кластеризация изображений - выделение памяти на GPU

Question

Я написал этот код для классификации изображений с помощью предварительно обученного googl enet:

gnet = models.googlenet(pretrained=True).cuda()

transform = transforms.Compose([transforms.Resize(256), transforms.CenterCrop(32), transforms.ToTensor()])
images = {}
resultDist = {}
i = 1

for f in glob.iglob("/data/home/student/HW3/trainData/train2014/*"):
    print(i)
    i = i + 1
    image = Image.open(f)
    # transform, create batch and get gnet weights
    img_t = transform(image).cuda()
    batch_t = torch.unsqueeze(img_t, 0).cuda()
    try:
        gnet.eval()
        out = gnet(batch_t)
        resultDist[f[-10:-4]] = out
        del out
    except:
        print(img_t.shape)
    del img_t
    del batch_t
    image.close()
    torch.cuda.empty_cache()
    i = i + 1

torch.save(resultDist, '/data/home/student/HW3/googlenetOutput1.pkl')

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

resultDist[f[-10:-4]] = out

Словарь занимает много места, и я не могу удалить его, потому что я хочу сохранить свои данные в файл pkl.

Answer 1

Поскольку вы не выполняете backprop, оберните весь l oop оператором with torch.no_grad():, так как в противном случае будет создан график вычислений и непостоянные результаты могут быть сохранены в GPU для последующего применения backprop. Это занимает изрядное количество места. Также вы, вероятно, хотите сохранить out.cpu(), чтобы ваши результаты не оставались на графическом процессоре.

...
with torch.no_grad():
    for f in glob.iglob("/data/home/student/HW3/trainData/train2014/*"):
        ...
            resultDist[f[-10:-4]] = out.cpu()
        ...

torch.save(resultDist, '/data/home/student/HW3/googlenetOutput1.pkl')