Pytorch保存整个网络和保存网络参数

Pytorch保存整个网络和保存网络参数

两种方法

---

定义网络和训练：

 net1 = torch.nn.Sequential(
        torch.nn.Linear(1, 10),
        torch.nn.ReLU(),
        torch.nn.Linear(10, 1)
    )
 optimizer = torch.optim.SGD(net1.parameters(), lr=0.5)
 loss_func = torch.nn.MSELoss()
 # 训练
 for t in range(100):
     prediction = net1(x)
     loss = loss_func(prediction, y)
     optimizer.zero_grad()
     loss.backward()
     optimizer.step()

---

两种保存的方法：

torch.save(net1, 'net.pkl')  # 保存整个网络
torch.save(net1.state_dict(), 'net_params.pkl')   # 只保存网络中的参数 (速度快, 占内存少)

提取网络：

# 复制整个网络
net2 = torch.load('net.pkl')  

# 仅提取网络参数
# 新建 net3
net3 = torch.nn.Sequential(
       torch.nn.Linear(1, 10),
       torch.nn.ReLU(),
       torch.nn.Linear(10, 1)
    )

# 将保存的参数复制到 net3
net3.load_state_dict(torch.load('net_params.pkl'))

需要注意的事：

当你复制整个网络，代码中要有网络的类，否则会报错，具体原因我猜测应该是pytorch中有一个自动实例化网络类的过程

---