从预先训练的模型中获取一组嵌入-vgg16 pytorch

Question

出于某个项目的目的，我试图将大约6000张图像的1 * 4096嵌入（最后一层之前的输出）存储到pkl文件中。同样，我在Google colab的vgg16修改模型上对6000张图像进行了迭代。但它会在内存中返回“ CUDA”。尝试分配14.00 MiB（GPU 0; 15.90 GiB总容量; 14.86 GiB已分配; 1.88 MiB可用; 342.26 MiB已缓存）'错误。

尽管我使用了相同的数据集（分为测试序列）来训练和验证我的模型，但运行良好。我想知道为什么单独获取和存储嵌入在colab中成为一项繁重的任务。 除了下面的代码，还有其他方法可以获取嵌入内容并将其存储在pkl文件中吗？

embedding = []
vgg16 = vgg16.to(device)
for x in range (0, len(inputImages)) :
  input = transformations(inputImages[x]) //pre processing
  input = torch.unsqueeze(input, 0)
  input = input.to(device)
  embedding.append(vgg16(input))

代码在最后一行与CUDA内存不足错误中断。

Answer 1

您生成的vgg16(input)输出仍然在cuda中。之所以如此，是因为此输出用于随后计算损耗。因此，为避免将输出存储在CUDA中并耗尽您的GPU内存，请使用.cpu().numpy()将其移至CPU。如果抛出错误，则可能还必须使用.detach()来分离变量。