图像聚类 - 在 GPU 上分配内存

Question

我已经为预训练的 googlenet 的图像分类编写了这段代码：

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')

我在使用它们后从 GPU 中删除了所有可能的张量，但是在我的数据集中大约 8000 张图像之后，GPU 已满。我发现问题出在：

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

字典占用了很多space，我无法删除它，因为我想将我的数据保存到pkl文件中。

Answer 1

因为你没有做反向传播，所以用 with torch.no_grad(): 语句包装你的整个循环，否则会创建一个计算图，并且间歇性的结果可能会存储在 GPU 上以供以后应用反向传播。这需要相当多的 space。此外，您可能想要保存 out.cpu()，这样您的结果就不会留在 GPU 上。

...
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')

图像聚类 - 在 GPU 上分配内存

Image clustering - allocating memory on GPU

python

image

classification

pytorch