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()
  1. 然后保存模型参数,共有两种方法
# 方法一:保存整个网络以及参数
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'))

你可能感兴趣的:(pytorch,pytorch,深度学习,python)