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

字典占用了大量空间，我无法删除它，因为我想将数据保存到pkl文件中。

Answer 1

由于您没有进行反向传播，因此用with torch.no_grad():语句包装了整个循环，因为否则会创建计算图，并且间歇性结果可能会存储在GPU上，以便以后应用反向传播。这需要相当大的空间。另外，您可能想保存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上分配内存

1 个答案: