我正在尝试构建一个VGG16模型以使用PyTorch导出ONNX。我希望用我自己的权重和偏置集来强制模型。在这个过程中，我的电脑很快就耗尽了内存。
这是我想做的方式（这只是一个测试，在实际版本中，我会从一组文件中读取权重和偏置），这个例子只是将所有值强制设为0.5
# 创建空的VGG16模型（随机权重）from torchvision import modelsfrom torchsummary import summaryvgg16 = models.vgg16()# la structure est : vgg16.__dict__summary(vgg16, (3, 224, 224))# 卷积层for layer in vgg16.features: print() print(layer) if (hasattr(layer,'weight')): dim = layer.weight.shape print(dim) print(str(dim[0]*(dim[1]*dim[2]*dim[3]+1))+' params') # 替换权重和偏置 for i in range (dim[0]): layer.bias[i] = 0.5 for j in range (dim[1]): for k in range (dim[2]): for l in range (dim[3]): layer.weight[i][j][k][l] = 0.5# 全连接层for layer in vgg16.classifier: print() print(layer) if (hasattr(layer,'weight')): dim = layer.weight.shape print(str(dim)+' --> '+str(dim[0]*(dim[1]+1))+' params') for i in range(dim[0]): layer.bias[i] = 0.5 for j in range(dim[1]): layer.weight[i][j] = 0.5当我查看电脑的内存使用情况时，它会线性增长，并在处理第一个全连接层时饱和了16GB的RAM。然后Python崩溃了…
有没有更好的方法来做这件事，考虑到之后我希望导出ONNX模型？感谢您的帮助。
回答：
内存增长的原因是需要为每次权重和偏置的更改调整梯度。尝试在更新前将.requires_grad属性设置为False，并在更新后恢复它。例如：
for layer in vgg16.features: print() print(layer) if (hasattr(layer,'weight')): # 抑制 .requires_grad layer.bias.requires_grad = False layer.weight.requires_grad = False dim = layer.weight.shape print(dim) print(str(dim[0]*(dim[1]*dim[2]*dim[3]+1))+' params') # 替换权重和偏置 for i in range (dim[0]): layer.bias[i] = 0.5 for j in range (dim[1]): for k in range (dim[2]): for l in range (dim[3]): layer.weight[i][j][k][l] = 0.5 # 恢复 .requires_grad layer.bias.requires_grad = True layer.weight.requires_grad = True