pytorch 模型的保存与加载

首先，创建一个简单的网络模型：

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x
net = MyModel()

2. 然后保存模型参数，共有两种方法

# 方法一：保存整个网络以及参数
torch.save(net,'net.pkl')


# 方法二：只保存网络中的参数（速度快，占用内存少），但是在加载时需要自己实现网络结构
torch.save(net.state_dict(), 'net_params.pkl')

当然也可以保存为.pt, .pkl, .rar，只需要更改后缀即可，如果没有特殊需求，使用任何一种都可以实现模型的保存和加载。
3. 最后加载保存的模型到新网络中，分别对应上述两种保存方法

# 对应保存方法一
new_net = torch.load('net.pkl')


# 对应保存方法二
new_net = torch.nn.Sequential( # 首先要实现对应参数的网络模型，然后才能加载参数
	torch.nn.Conv2d(3, 6, 5),
	torch.nn.MaxPool2d(2, 2),
	torch.nn.Conv2d(6, 16, 5),
	torch.nn.Linear(16 * 5 * 5, 120),
	torch.nn.Linear(120, 84),
	torch.nn.Linear(84, 10))

new_net.load_state_dict(torch.load('net_params.pkl'))