从零搭建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_()

搭建网络有几个要点：

1. 自定义类要继承torch.nn.Module。有时候自己设计了一些模块，为了使用更方便，通常额外定义一个类，就像这里的Flatten，自定义的类也要继承torch.nn.Module。

2. 完成init函数和forward函数。其中__init__函数完成网络的搭建，forward函数完成网络的前传路径。

3. 完成所有层的参数初始化，一般只有卷积层，归一化层，全连接层要初始化，池化层没有参数。

__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 out

2.下面是设置不同的层。注：上面和下面都不是一个完整的代码，只是用来说明这种很多层的构建方式。

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数据增强策略