在PyTorch里面写神经网络,所有的层结构和损失函数都来自于torch.nn,所有的模型构建都是从nn.Module这个基类继承的。
模型的创建过程就是定义一个class继承nn.Module,然后重写init方法和forward方法,实例代码如下:
import torch.nn as nn
import torch.nn.functional as F
class Model(nn.Module):
def __init__(self):
super(Model, self).__init__()
self.conv1 = nn.Conv2d(1, 20, 5)
self.conv2 = nn.Conv2d(20, 20, 5)
def forward(self, x):
x = F.relu(self.conv1(x))
return F.relu(self.conv2(x))
定义好模型后,需要通过nn这个包来定义损失函数。
常见的损失函数都已经定义在nn中,如均方差、多分类交叉熵
criterion = nn.CrossEntropyLoss()
loss = criterion(output, target)
需要修改参数使损失函数最小化
torch.optim是一个实现各种优化算法的包
model = net_name()
optimizer = torch.optim.SGD(model.parameters(),lr=0.01,momentum=0.9)
在优化之前需要先将梯度归零,然后进行反向传播
optimizer.zeros()
optimizer.step()
模型的保存和加载
模型的保存有两种方式,一种是保存整个模型的网络结构和参数,另一种是指保存模型节点的参数:
torch.save(model,'./model.pth')
torch.save(model.state_dict(),'./model_state.pth')
与之相对应,模型的加载也有两种方式
load_model = torch.load('./model.pth')
load_model_ = model.load_state_dict(torch.load('./model_state.pth'))