PyTorch模型定义
PyTorch模型定义应包括两个主要部分：
- 各个部分的初始化（_init_）；
- 数据流向定义（forward）。
基于nn.Module，可以通过Sequential，ModuleList和ModuleDict三种方式定义PyTorch模型。
方式1-Sequential
nn.Sequential()
接收一系列子模块或者一个子模块的有序字典(OrderedDict) 作为参数来逐一添加 Module 的实例。
使用Sequential定义的模型可以不重写forward
适用场景：
模型的前向传播(forward)为简单串联多个层的计算。
定义方式的理解：
Sequetial按照以下逻辑进行建模
# Sequential 的定义方式
class MySequential(nn.Module):
def __init__(self, *args):
super(MySequential, self).__init__()
if len(args) == 1 and isinstance(args[0], OrderedDict): # 如果传入的是一个OrderedDict
for key, module in args[0].items():
self.add_module(key, module) # add_module方法会将module添加进self._modules(一个OrderedDict)
else: # 传入的是一些Module
for idx, module in enumerate(args):
self.add_module(str(idx), module)
def forward(self, input):
# self._modules返回一个 OrderedDict，保证会按照成员添加时的顺序遍历成
for module in self._modules.values():
input = module(input)
return input
使用方法：
直接排列
model1 = nn.Sequential(
nn.Linear(784, 256),
nn.ReLU(),
nn.Linear(256, 10),
)
print(model1)
使用OrderedDict
model2 = nn.Sequential(OrderedDict({
'layer1': nn.Linear(784,256),
'relu1': nn.ReLU(),
'layer2': nn.Linear(256,10)}))
print(model2)
方式2-ModuleList
nn.ModuleList()
接收一个子模块（或层，需属于nn.Module类）的列表(list)作为输入，可使用类似List的操作，不拿包括append(module)、extend(modules)、insert(index, module)。同时，子模块或层的权重也会自动添加到网络中来。
适用场景：
在某个完全相同的层需要重复出现多次时
使用方法：
model3 = nn.ModuleList([nn.Linear(784,256), nn.ReLU()])
model3.append(nn.Linear(256, 10)) # 类似List的append操作
print(model3[-1]) # 类似List的索引访问
print(model3)
使用ModuleList定义的模型必须重写forward，用for循环即可：
class My_Model(nn.Module):
def __init__(self, input_dim):
super(My_Model, self).__init__()
self.modulelist = nn.ModuleList([nn.Linear(input_dim,256),
nn.ReLU(),
nn.Linear(256, 10)])
def forward(self, x):
for layer in self.modulelist:
x = layer(x)
return x
方式3-ModuleDict
nn.ModuleDict()
和ModuleList的作用类似，但能够更方便地为神经网络的层添加名称。
适用场景：
在某个完全相同的层需要重复出现多次时
使用方法：
model4 = nn.ModuleDict({
'linear': nn.Linear(784, 256),
'act': nn.ReLU(),
})
model4['output'] = nn.Linear(256, 10) # 添加
print(model4['linear']) # 访问
print(model4.output)
print(model4)
组成U-Net的模型块主要有如下几个部分：
1）每个子块内部的两次卷积（Double Convolution）—— DoubleConv
2）左侧模型块之间的下采样连接，即最大池化（Max pooling）—— Down
3）右侧模型块之间的上采样连接（Up sampling）—— Up
4）输出层的处理 —— OutConv
除模型块外，还有模型块之间的横向连接，输入和U-Net底部的连接等计算，这些单独的操作可以通过forward函数来实现。
具体代码参考
class DoubleConv(nn.Module):
"""
(convolution => [BN] => ReLU) * 2
in_channels (int): Number of channels in the input image
out_channels (int): Number of channels produced by the convolution
kernel_size (int or tuple): Size of the convolving kernel
"""
def __init__(self, in_channels, out_channels, mid_channels=None):
super().__init__()
if not mid_channels:
mid_channels = out_channels
self.double_conv = nn.Sequential(
nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1, bias=False),
nn.BatchNorm2d(mid_channels),
nn.ReLU(inplace=True),
nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1, bias=False),
nn.BatchNorm2d(out_channels),
nn.ReLU(inplace=True)
)
def forward(self, x):
return self.double_conv(x)
class Down(nn.Module):
"""Downscaling with maxpool then double conv"""
def __init__(self, in_channels, out_channels):
super().__init__()
self.maxpool_conv = nn.Sequential(
nn.MaxPool2d(2),
DoubleConv(in_channels, out_channels)
)
def forward(self, x):
return self.maxpool_conv(x)
class Up(nn.Module):
"""Upscaling then double conv"""
def __init__(self, in_channels, out_channels, bilinear=True):
super().__init__()
# if bilinear, use the normal convolutions to reduce the number of channels
if bilinear:
self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
self.conv = DoubleConv(in_channels, out_channels, in_channels // 2)
else:
self.up = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2)
self.conv = DoubleConv(in_channels, out_channels)
def forward(self, x1, x2):
x1 = self.up(x1)
# input is CHW
diffY = x2.size()[2] - x1.size()[2]
diffX = x2.size()[3] - x1.size()[3]
x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2,
diffY // 2, diffY - diffY // 2])
# if you have padding issues, see
# https://github.com/HaiyongJiang/U-Net-Pytorch-Unstructured-Buggy/commit/0e854509c2cea854e247a9c615f175f76fbb2e3a
# https://github.com/xiaopeng-liao/Pytorch-UNet/commit/8ebac70e633bac59fc22bb5195e513d5832fb3bd
x = torch.cat([x2, x1], dim=1)
return self.conv(x)
class OutConv(nn.Module):
def __init__(self, in_channels, out_channels):
super(OutConv, self).__init__()
self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)
def forward(self, x):
return self.conv(x)
class UNet(nn.Module):
def __init__(self, n_channels, n_classes, bilinear=True):
super(UNet, self).__init__()
self.n_channels = n_channels
self.n_classes = n_classes
self.bilinear = bilinear
self.inc = DoubleConv(n_channels, 64)
self.down1 = Down(64, 128)
self.down2 = Down(128, 256)
self.down3 = Down(256, 512)
factor = 2 if bilinear else 1
self.down4 = Down(512, 1024 // factor)
self.up1 = Up(1024, 512 // factor, bilinear)
self.up2 = Up(512, 256 // factor, bilinear)
self.up3 = Up(256, 128 // factor, bilinear)
self.up4 = Up(128, 64, bilinear)
self.outc = OutConv(64, n_classes)
def forward(self, x):
x1 = self.inc(x)
x2 = self.down1(x1)
x3 = self.down2(x2)
x4 = self.down3(x3)
x5 = self.down4(x4)
x = self.up1(x5, x4)
x = self.up2(x, x3)
x = self.up3(x, x2)
x = self.up4(x, x1)
logits = self.outc(x)
return logits
以pytorch官方视觉库torchvision预定义好的模型ResNet50为例
修改模型若干层
Target：用ResNet50模型去做一个10分类的问题，并再加一层全连接层。
Method：修改模型的fc层，将其输出节点数替换为10。
- 导入已有模型
import torchvision.models as models
net = models.resnet50()
- 查看具体模型，找到需要修改的层
print(net)
假设要修改最后一层(fc)
(fc): Linear(in_features=2048, out_features=1000, bias=True)
- 定义自己要替换的模型
from collections import OrderedDict
classifier = nn.Sequential(OrderedDict([('fc1', nn.Linear(2048, 128)),
('relu1', nn.ReLU()),
('dropout1',nn.Dropout(0.5)),
('fc2', nn.Linear(128, 10)),
('output', nn.Softmax(dim=1))
]))
- 对要修改的层赋值
net.fc = classifier
print(net)
最后模型为：
添加额外输入
Target：利用已有的模型结构，在倒数第二层增加一个额外的输入变量add_variable来辅助预测。
Method：将原模型添加输入位置前的部分作为一个整体，同时在forward中定义好原模型不变的部分、添加的输入和后续层之间的连接关系，从而完成模型的修改。
class Model(nn.Module):
def __init__(self, net):
super(Model, self).__init__()
self.net = net
self.relu = nn.ReLU()
self.dropout = nn.Dropout(0.5)
self.fc_add = nn.Linear(1001, 10, bias=True)
self.output = nn.Softmax(dim=1)
def forward(self, x, add_variable):
x = self.net(x)
x = torch.cat((self.dropout(self.relu(x)), add_variable.unsqueeze(1)),1)
x = self.fc_add(x)
x = self.output(x)
return x
net1 = models.resnet50()
model = Model(net1).cpu()
torch.cat(tensors, dim=0, *, out=None) → Tensor
tensors (sequence of Tensors) – 任何相同类型的张量的python序列。提供的非空张量必须具有相同的形状，除所 cat 维度外的维数必须匹配。
x = torch.randn(2, 3)
y = torch.randn(1, 3)
z = torch.randn(1, 3)
torch.cat((x, y, z), 0)
ResNet50模型的后两层输入输出维度分别为：
(layer4): (1024, 2048)
(fc): (2048, 1000)
对外部输入变量"add_variable"进行unsqueeze操作是为了和net输出的tensor保持维度一致
训练中在输入数据的时候要给两个inputs.
outputs = model(inputs, add_var)
添加额外输出
Target: 输出1000维的倒数第二层和10维的最后一层结果。
Method: 修改模型定义中forward函数的return变量.
class Model(nn.Module):
def __init__(self, net):
super(Model, self).__init__()
self.net = net
self.relu = nn.ReLU()
self.dropout = nn.Dropout(0.5)
self.fc1 = nn.Linear(1000, 10, bias=True)
self.output = nn.Softmax(dim=1)
def forward(self, x, add_variable):
x1000 = self.net(x)
x10 = self.dropout(self.relu(x1000))
x10 = self.fc1(x10)
x10 = self.output(x10)
return x10, x1000
net = models.resnet50()
model = Model(net).cpu()
训练中在输入数据后会有两个outputs
out10, out1000 = model(inputs, add_var)
PyTorch的模型的存储格式
三种格式：
pkl
pt
pth
PyTorch如何存储模型
一个PyTorch模型主要包含两个部分：模型结构和权重。
其中模型是继承nn.Module的类，权重的数据结构是一个字典（key是层名，value是权重向量）。
存储也由此分为两种形式：存储整个模型（包括结构和权重），和只存储模型权重。
存储整个模型：
torch.save(model, save_dir)
只存储模型权重：
torch.save(model.state_dict, save_dir)
单卡与多卡训练下模型的保存与加载方法
如果要使用多卡训练的话，需要对模型使用torch.nn.DataParallel。
多卡并行的模型每层的名称前多了一个“module”。
由此，总结PyTorch框架下单卡/多卡下模型的保存和加载问题
单卡保存+单卡加载
# 保存+读取整个模型
torch.save(model, save_dir)
loaded_model = torch.load(save_dir)
loaded_model.cuda()
# 保存+读取模型权重
torch.save(model.state_dict(), save_dir)
loaded_dict = torch.load(save_dir)
loaded_model = models.resnet152() #注意这里需要对模型结构有定义
loaded_model.state_dict = loaded_dict
loaded_model.cuda()
单卡保存+多卡加载
多卡保存+单卡加载
多卡保存+多卡加载
参考：https://github.com/datawhalechina/thorough-pytorch