从零搭建Pytorch模型教程（二）搭建网络

前言

上一篇《从零搭建Pytorch模型教程（一）数据读取》中介绍了classdataset的几个要点，由哪些部分组成，每个部分需要完成哪些事情，如何进行数据增强，如何实现自己设计的数据增强。然后，介绍了分布式训练的数据加载方式，数据读取的整个流程，当面对超大数据集时，内存不足的改进思路。

本文介绍了如何搭建神经网络，构建网络的几种方式，前向传播的过程，几种初始化方式，如何加载预训练模型的指定层等内容。本文以CNN为例，下一篇介绍如何搭建Transformer网络。

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