从零搭建Pytorch模型教程（二）搭建网络
前言
上一篇《从零搭建Pytorch模型教程（一）数据读取》中介绍了classdataset的几个要点，由哪些部分组成，每个部分需要完成哪些事情，如何进行数据增强，如何实现自己设计的数据增强。然后，介绍了分布式训练的数据加载方式，数据读取的整个流程，当面对超大数据集时，内存不足的改进思路。
本文介绍了如何搭建神经网络，构建网络的几种方式，前向传播的过程，几种初始化方式，如何加载预训练模型的指定层等内容。本文以CNN为例，下一篇介绍如何搭建Transformer网络。
本文来自公众号CV技术指南的技术总结系列
欢迎关注公众号CV技术指南，专注于计算机视觉的技术总结、最新技术跟踪、经典论文解读、CV招聘信息。
搭建CNN网络
首先来看一个CNN网络 (以YOLO_v1的一部分层为例)。
class Flatten(nn.Module):
def __init__(self):
super(Flatten,self).__init__()
def forward(self,x):
return x.view(x.size(0),-1)
class Yolo_v1(nn.Module):
def __init__(self, num_class):
super(Yolo_v1,self).__init__()
C = num_class
self.conv_layer1=nn.Sequential(
nn.Conv2d(in_channels=3,out_channels=64,kernel_size=7,stride=1,padding=7//2),
nn.BatchNorm2d(64),
nn.LeakyReLU(0.1),
nn.MaxPool2d(kernel_size=2,stride=2)
)
self.conv_layer2=nn.Sequential(
nn.Conv2d(in_channels=64,out_channels=192,kernel_size=3,stride=1,padding=3//2),
nn.BatchNorm2d(192),
nn.LeakyReLU(0.1),
nn.MaxPool2d(kernel_size=2,stride=2)
)
#为了简便，这里省去了很多层
self.flatten = Flatten()
self.conn_layer1 = nn.Sequential(
nn.Linear(in_features=7*7*1024,out_features=4096),
nn.Dropout(0.5),nn.LeakyReLU(0.1))
self.conn_layer2 = nn.Sequential(nn.Linear(in_features=4096,out_features=7*7*(2*5 + C)))
self._initialize_weights()
def forward(self,input):
conv_layer1 = self.conv_layer1(input)
conv_layer2 = self.conv_layer2(conv_layer1)
flatten = self.flatten(conv_layer2)
conn_layer1 = self.conn_layer1(flatten)
output = self.conn_layer2(conn_layer1)
return output
def _initialize_weights(self):
for m in self.modules():
if isinstance(m, nn.Conv2d):
n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
m.weight.data.normal_(0, math.sqrt(2. / n))
if m.bias is not None:
m.bias.data.zero_()
elif isinstance(m, nn.BatchNorm2d):
m.weight.data.fill_(1)
m.bias.data.zero_()
elif isinstance(m, nn.Linear):
m.weight.data.normal_(0, 0.01)
m.bias.data.zero_()
搭建网络有几个要点：
-
自定义类要继承torch.nn.Module。有时候自己设计了一些模块，为了使用更方便，通常额外定义一个类，就像这里的Flatten，自定义的类也要继承torch.nn.Module。
-
完成init函数和forward函数。其中__init__函数完成网络的搭建，forward函数完成网络的前传路径。
-
完成所有层的参数初始化，一般只有卷积层，归一化层，全连接层要初始化，池化层没有参数。
__init__函数
构建网络层有几种方式，一种是pytorch官方已经有了定义的网络，如resnet，vgg，Inception等。一种是自定义层，例如自己设计了一个新的模块。
首先是使用pytorch官方库已经支持的网络，这些网络放在了torchvision.models中,下面选择自己需要的一个。
以下只列举了2D 模型的一部分，还有视频类的3D 模型。
import torchvision.models as models
resnet18 = models.resnet18(pretrained = True)
alexnet = models.alexnet()
vgg16 = models.vgg16()
squeezenet = models.squeezenet1_0()
densenet = models.densenet161()
inception = models.inception_v3()
googlenet = models.googlenet()
shufflenet = models.shufflenet_v2_x1_0()
mobilenet_v2 = models.mobilenet_v2()
mobilenet_v3_large = models.mobilenet_v3_large()
mobilenet_v3_small = models.mobilenet_v3_small()
resnext50_32x4d = models.resnext50_32x4d()
wide_resnet50_2 = models.wide_resnet50_2()
mnasnet = models.mnasnet1_0()
efficientnet_b0 = models.efficientnet_b0()
efficientnet_b1 = models.efficientnet_b1()
efficientnet_b2 = models.efficientnet_b2()
regnet_y_400mf = models.regnet_y_400mf()
regnet_y_800mf = models.regnet_y_800mf()
vit_b_16 = models.vit_b_16()
vit_b_32 = models.vit_b_32()
vit_l_16 = models.vit_l_16()
vit_l_32 = models.vit_l_32()
convnext_tiny = models.convnext_tiny()
convnext_small = models.convnext_small()
convnext_base = models.convnext_base()
convnext_large = models.convnext_large()
若需要加载该网络在ImageNet上预训练的模型，则在括号内设置参数pretrained=True即可。但这种方式有个不好的问题在于这些预训练模型并不是在本地，因此每次运行都会从网上读取加载模型，非常浪费时间。因此，可以去它官网(https://pytorch.org/)上把那个模型下载到本地，通过下面指令完成加载。
resnet50.load_state_dict(torch.load('/path/to/resnet50.pth'))另一种自定义层的，一般可以通过torch.nn.Sequential()来构建，在中间插入卷积层、归一化层、激活函数层、池化层即可。
例如下方这种是最常用的。
self.conv_layer1=nn.Sequential(
nn.Conv2d(in_channels=3,out_channels=64,kernel_size=7,stride=1,padding=7//2),
nn.BatchNorm2d(64),
nn.LeakyReLU(0.1),
nn.MaxPool2d(kernel_size=2,stride=2)，
nn.Conv2d(in_channels=3,out_channels=64,kernel_size=7,stride=1,padding=7//2),
nn.BatchNorm2d(64),
nn.LeakyReLU(0.1),
nn.MaxPool2d(kernel_size=2,stride=2)，
)
当网络很深时，上面这种方式构建比较麻烦，例如resnet，总不可能就按找上面这种方式这么写50层。就把它们共同的部分给构建出来，然后通过传参来设置不同的层。
例如：
1.下面这里先构建一个基本的几层作为一个类，每一层的参数(不同输入输出通道数，卷积核大小，有无池化)都通过传参来设置。
class BasicBlock(nn.Module):
expansion = 1
def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
base_width=64, dilation=1, norm_layer=None):
super(BasicBlock, self).__init__()
self.conv1 = conv3x3(inplanes, planes, stride)
self.bn1 = norm_layer(planes)
self.relu = nn.ReLU(inplace=True)
self.conv2 = conv3x3(planes, planes)
self.bn2 = norm_layer(planes)
self.downsample = downsample
self.stride = stride
def forward(self, x):
identity = x
out = self.conv1(x)
out = self.bn1(out)
out = self.relu(out)
out = self.conv2(out)
out = self.bn2(out)
if self.downsample is not None:
identity = self.downsample(x)
out += identity
out = self.relu(out)
return out2.下面是设置不同的层。注：上面和下面都不是一个完整的代码，只是用来说明这种很多层的构建方式。
layers = []
layers.append(block(self.inplanes, planes, stride, downsample, self.groups,
self.base_width, previous_dilation, norm_layer))
self.inplanes = planes * block.expansion
for _ in range(1, blocks):
layers.append(block(self.inplanes, planes, groups=self.groups,
base_width=self.base_width, dilation=self.dilation,
norm_layer=norm_layer))
return nn.Sequential(*layers)
forward函数
这里就是网络的传播路径了，一般就是一路往下传就是。return的内容就是网络的输出。
def forward(self,x):
x = self.conv_layer1(x)
x = self.conv_layer2(x)
x = sexlf.flatten(x)
x = self.conn_layer1(x)
output = self.conn_layer2(x)
return output
如果想将中间某几层的输出拿出来，做一下特征金字塔，可以像下面这么写。
def forward(self,x):
conv_layer1 = self.conv_layer1(x)
conv_layer2 = self.conv_layer2(conv_layer1)
conv_layer3 = self.conv_layer2(conv_layer2)
conv_layer4 = self.conv_layer2(conv_layer3)
FP = self.YourModule(conv_layer1,conv_layer2,conv_layer3,conv_layer4)
flatten = self.flatten(FN)
conn_layer1 = self.conn_layer1(flatten)
output = self.conn_layer2(conn_layer1)
return output
像可视化特征图里，想要可视化某一层的特征图，就可以像下面这么写。
def forward(self,x):
x = self.conv_layer1(x)
feature = self.conv_layer2(x)
x = sexlf.flatten(feature)
x = self.conn_layer1(x)
output = self.conn_layer2(x)
return feature,output
初始化网络
初始化网络是要放在init函数里完成，分为两类，一类是随机初始化，一类是加载预训练模型。
随机初始化
关于随机初始化，目前主要有多种方式：Normal Initialization, Uniform Initialization，Xavier Initialization，He Initialization (也称 kaiming Initialization)，LeCun Initialization。
关于这些初始化方法，可以看这篇文章《神经网络的初始化方法总结 | 又名“如何选择合适的初始化方法”》。我们一般使用Kaiming Initialization。
下面是一种方式，直接按自定义的方式初始化。
def _initialize_weights(self):
for m in self.modules():
if isinstance(m, nn.Conv2d):
n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
m.weight.data.normal_(0, math.sqrt(2. / n))
if m.bias is not None:
m.bias.data.zero_()
elif isinstance(m, nn.BatchNorm2d):
m.weight.data.fill_(1)
m.bias.data.zero_()
elif isinstance(m, nn.Linear):
m.weight.data.normal_(0, 0.01)
m.bias.data.zero_()
也可以选择pytorch实现了的初始化。
for m in self.modules():
if isinstance(m, nn.Conv2d):
nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
nn.init.constant_(m.weight, 1)
nn.init.constant_(m.bias, 0)
还可以像下面这么写：
from torch.nn import init
def weights_init_kaiming(m):
classname = m.__class__.__name__
# print(classname)
if classname.find('Conv') != -1:
init.kaiming_normal_(m.weight.data, a=0, mode='fan_in') # For old pytorch, you may use kaiming_normal.
elif classname.find('Linear') != -1:
init.kaiming_normal_(m.weight.data, a=0, mode='fan_out')
init.constant_(m.bias.data, 0.0)
elif classname.find('BatchNorm2d') != -1:
init.normal_(m.weight.data, 1.0, 0.02)
init.constant_(m.bias.data, 0.0)
def weights_init_classifier(m):
classname = m.__class__.__name__
if classname.find('Linear') != -1:
init.normal_(m.weight.data, std=0.001)
init.constant_(m.bias.data, 0.0)
self.conv_layer1=nn.Sequential(
nn.Conv2d(in_channels=3,out_channels=64,kernel_size=7,stride=1,padding=7//2),
nn.BatchNorm2d(64),
nn.LeakyReLU(0.1),
nn.MaxPool2d(kernel_size=2,stride=2)
)
self.conv_layer1.apply(weights_init_kaiming)
反正随便选择一种就好。
加载预训练模型初始化
加载预训练模型一般是在train文件里写，但有些网络由于是使用现成的backbone网络，例如使用了resnet50，然后后面加了自定义的模块，所以它想要resnet50预训练模型初始化backbone，而其它层做随机初始化，那加载预训练模型就是在网络定义中做的。因此，既然这里提到了初始化，就干脆写在这里。
最简单的就是直接整个模型都加载。
resnet50.load_state_dict(torch.load('/path/to/resnet50.pth'))
但也有一些情况下，我只想加载其中一部分层的参数。剩下一部分由于已经改变参数了，无法加载预训练模型，所以要选择上面的随机初始化。
这里有必要来说明网络的每一层是如何表示的。下面以一个例子来说明。
class Flatten(nn.Module):
def __init__(self):
super(Flatten,self).__init__()
def forward(self,x):
return x.view(x.size(0),-1)
class YourNet(nn.Module):
def __init__(self,stride=2, pool='avg'):
super(YourNet, self).__init__()
self.resnet50 = models.resnet50(pretrained=False)
self.model.load_state_dict(torch.load('/path/to/resnet50.pth'))
self.flatten = Flatten()
self.conn_layer1 = nn.Sequential(
nn.Linear(in_features=7 * 7 * 1024, out_features=4096),
nn.Dropout(0.5),
nn.LeakyReLU(0.1)
)
self.conn_layer2 = nn.Sequential(nn.Linear(in_features=4096, out_features=7 * 7 * (2 * 5 + 20)))
def forward(self,x):
#这里省略
if __name__ == "__main__":
model = YourNet()
for name, value in model.named_parameters():
print(name)
这里简单定义了一个网络。在最后面有这两行：
for name, value in model.named_parameters():
print(name)
这两行的输出就是打印网络层的名字，实际上加载预训练模型时，也是按照这个名字来加载的。下面是一部分输出。
resnet50.conv1.weight
resnet50.bn1.weight
resnet50.bn1.bias
resnet50.layer1.0.conv1.weight
resnet50.layer1.0.bn1.weight
resnet50.layer1.0.bn1.bias
resnet50.layer1.0.conv2.weight
resnet50.layer1.0.bn2.weight
resnet50.layer1.0.bn2.bias
resnet50.layer1.0.conv3.weight
resnet50.layer1.0.bn3.weight
resnet50.layer1.0.bn3.bias
resnet50.layer1.0.downsample.0.weight
resnet50.layer1.0.downsample.1.weight
resnet50.layer1.0.downsample.1.bias
...
...
resnet50.layer4.2.bn3.weight
resnet50.layer4.2.bn3.bias
resnet50.fc.weight
resnet50.fc.bias
conn_layer1.0.weight
conn_layer1.0.bias
conn_layer2.0.weight
conn_layer2.0.bias
在预训练模型中就是这样，key即为网络层的名字，value即为它们对应的参数。因此，加载预训练模型可以按照下面这种方式加载。
pretrained_dict = torch.load('/path/to/resnet50.pth')
pretrained_dict.pop('fc.weight')
pretrained_dict.pop('fc.bias')
#自己的模型参数变量
model_dict = model.state_dict()
#去除一些不需要的参数
pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
#参数更新
model_dict.update(pretrained_dict)
# 加载我们真正需要的state_dict
model.load_state_dict(model_dict)
自己定义的一些层是不会出现在pretrained_dict中，因此会将其剔除，从而只加载了pretrained_dict中有的层。
本文介绍了如何搭建神经网络，构建网络的几种方式，前向传播的过程，几种初始化方式，如何加载预训练模型的指定层等内容。
下一篇我们将介绍如何写train函数，以及包括设置优化方式，设置学习率，不同层设置不同学习率，解析参数等。
欢迎关注公众号CV技术指南，专注于计算机视觉的技术总结、最新技术跟踪、经典论文解读、CV招聘信息。
CV技术指南创建了一个交流氛围很不错的群，除了太偏僻的问题，几乎有问必答。关注公众号添加编辑的微信号可邀请加交流群。
在公众号中回复关键字 “入门指南“可获取计算机视觉入门所有必备资料。
其它文章
用什么tricks能让模型训练得更快？模型训练慢的可能原因总结
经典论文系列 | 缩小Anchor-based和Anchor-free检测之间差距的方法：自适应训练样本选择
小样本学习研究综述（中科院计算所）
目标检测中正负样本区分策略和平衡策略总结
模型量化技巧及在低功耗IOT设备的应用实践
一份可视化特征图的代码
NeurIPS 2021 | 从进化算法角度解释Transformer架构，并提出针对多模态任务的统一序列模型范式
CVPR2022 | 单GPU每秒76帧，重叠对象也能完美分割，多模态Transformer用于视频分割效果惊艳
目标检测中的框位置优化总结
目标检测、实例分割、多目标跟踪的Anchor-free应用方法总结
Soft Sampling：探索更有效的采样策略
通用文字识别的问题和算法
PRCV 2021 | 腾讯优图分享内容理解新实践
如何解决工业缺陷检测小样本问题
CVPR2021 | TrivialAugment：不用调优的SOTA数据增强策略