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

【问题标题】：Obtain a set of embedding from pretrained model - vgg16 pytorch从预训练模型中获取一组嵌入 - vgg16 pytorch
【发布时间】：2020-01-03 06:54:48
【问题描述】：

出于某个项目目的，我试图将大约 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 内存不足错误。

【问题讨论】：

标签： python deep-learning conv-neural-network pytorch

【解决方案1】：

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

【讨论】：