Pytorch入门（二）用Pytorch写第一个神经网络

本文将用Pytorch构建一个最简单的线性神经网络，练习Pytorch中神经网络模型的保存和重载

定义一个线性神经网络

$$y=wx+b$$
这是一个基本的网络Net，它只包含一个全连接层

import torch
from torch import nn


class Net(nn.Module):
    def __init__(self):
        super(Net,self).__init__()
        self.layer = nn.Linear(1,1)
        self.layer.weight = nn.Parameter(torch.FloatTensor([[10]]))
        self.layer.bias = nn.Parameter(torch.FloatTensor([1]))

    def forward(self,x):
        y = self.layer(x)
        return y

x = torch.FloatTensor([[1]])
net = Net()
linearout = net(x)
print(linearout)

这里假设输入x=1
y的值应为11

tensor([[11.]], grad_fn=<AddmmBackward>)

保存Net的参数值

查看网络的状态字典

print(net.state_dict())

OrderedDict([('layer.weight', tensor([[10.]])), ('layer.bias', tensor([1.]))])

保存状态字典

torch.save(obj=net.state_dict(), f="models/net.pth")

加载Net参数值并用于新的模型

重新定义一个相同结构的模型NewNet

import torch
from torch import nn


class Net(nn.Module):
    def __init__(self):
        super(Net,self).__init__()
        self.layer = nn.Linear(1,1)
        self.layer.weight = nn.Parameter(torch.FloatTensor([[10]]))
        self.layer.bias = nn.Parameter(torch.FloatTensor([1]))

    def forward(self,x):
        y = self.layer(x)
        return y


class NewNet(nn.Module):
    def __init__(self):
        super(NewNet, self).__init__()
        self.layer = nn.Linear(1, 1)
        self.layer.weight = nn.Parameter(torch.FloatTensor([[0]]))
        self.layer.bias = nn.Parameter(torch.FloatTensor([0]))

    def forward(self, x):
        out = self.layer(x)
        return out


x = torch.FloatTensor([[1]])
net = NewNet()
print(net.state_dict())                             # 初始的NewNet的状态字典
net.load_state_dict(torch.load("models/net.pth"))
print(net.state_dict())                             # 加载参数值的NewNet的状态字典

net的w和b值就不再是0了，而是之前保存的模型中w和b对应的10和1

OrderedDict([('layer.weight', tensor([[0.]])), ('layer.bias', tensor([0.]))])
OrderedDict([('layer.weight', tensor([[10.]])), ('layer.bias', tensor([1.]))])

优化器与epoch的保存

保存优化器参数值和epoch值的主要目的是用于继续训练，保存的流程依旧是先“torch.save()”再“torch.load_state_dict()”
我们首先定义一个Adam优化器、一个任意的epoch值与net如下：

mport torch
from torch import nn,optim


class Net(nn.Module):
    def __init__(self):
        super(Net,self).__init__()
        self.layer = nn.Linear(1,1)
        self.layer.weight = nn.Parameter(torch.FloatTensor([[10]]))
        self.layer.bias = nn.Parameter(torch.FloatTensor([1]))

    def forward(self,x):
        y = self.layer(x)
        return y


net = Net()
Adam = optim.Adam(params=net.parameters(), lr=0.001, betas=(0.5, 0.999))
epoch = 50
all_states = {"net": net.state_dict(), "Adam": Adam.state_dict(), "epoch": epoch}
torch.save(obj=all_states, f="models/all_states.pth")

查看模型所有的参数

all_states = torch.load("models/all_states.pth")
print(all_states)

{
	'net': OrderedDict([('layer.weight', tensor([[10.]])), ('layer.bias', tensor([1.]))]),
	'Adam': {
		'state': {},
		'param_groups': [{
			'lr': 0.001,
			'betas': (0.5, 0.999),
			'eps': 1e-08,
			'weight_decay': 0,
			'amsgrad': False,
			'params': [4660776392, 4559888248]
		}]
	},
	'epoch': 50
}