从预训练的vgg16模型中得到一组嵌入

Question

为了某个项目的目的，我试图将大约6000个图像的1*4096嵌入（最后一层之前的输出）存储到一个pkl文件中。同样，我正在googlecolab中对vgg16修改模型上的6000个图像进行迭代。但它会从记忆中返回“CUDA”。尝试分配14.00 MiB（GPU 0；总容量15.90 GiB；14.86吉布已分配；1.88无MiB；342.26 MiB cached）'错误

然而，我使用了同样的数据集，并将其拆分为测试序列，用于训练和验证我的模型，运行良好。我想知道为什么单独获取和存储嵌入在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()来分离变量