首先,创建一个简单的网络模型:
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()
- 然后保存模型参数,共有两种方法
# 方法一:保存整个网络以及参数
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'))