保存和加载 pytorch 神经网络答案

【问题标题】：saving and loading pytorch neural nets保存和加载 pytorch 神经网络
【发布时间】：2019-05-16 14:53:18
【问题描述】：

所以我创建了一个神经网络，我想保存它并随时加载它。具体来说，我想拍照并进行实时处理。我正在使用创建的神经网络here

我读到标准方法是创建网络，然后使用torch.save(net,'mynet') 保存它，然后使用torch.load('mynet') 加载它。

但是，如果我打开一个新的python3 终端并使用：

>>import torch
>>torch.load('mynet')

它给了我错误：

  File "<stdin>", line 1, in <module>
  File "/home/tim/anaconda3/lib/python3.7/site-packages/torch/serialization.py", line 367, in load
    return _load(f, map_location, pickle_module)
  File "/home/tim/anaconda3/lib/python3.7/site-packages/torch/serialization.py", line 538, in _load
    result = unpickler.load()
AttributeError: Can't get attribute 'Net' on <module '__main__' (built-in)>

我认为这是因为没有定义 Net 类。添加

import torch.nn as nn
import torch.nn.functional as F
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 15, 3)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(15, 15, 5)
        self.conv3 = nn.Conv2d(15, 10, 3)
        self.fc1 = nn.Linear(10*4*4, 100)
        self.fc2 = nn.Linear(100, 24)
        self.fc3 = nn.Linear(24, 4)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = self.pool(F.relu(self.conv3(x)))
        x = x.view(-1, 10*4*4)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

是你需要做的，但是为什么我们需要定义神经网络类呢？如果我加载一个与我在类中指定的架构不同的神经网络，会覆盖类中定义的架构吗？我加载的对象肯定包含了所有的架构和类信息吗？

更新：实际上，当我定义 Net 类时它甚至都不起作用。

【问题讨论】：

标签： python-3.x image-processing pytorch

【解决方案1】：

请参阅serialization semantics 上的文档，该文档首先描述了建议的方法，然后是您用作“序列化数据绑定到特定类和使用的确切目录结构，因此它可以在各种情况下以各种方式中断”在其他项目中使用，或者经过一些严重的重构。”

换句话说，您需要保存/加载net.state_dict()，而不是net 本身。

【讨论】：