【PyTorch学习】（一）ResNet源码研读

       最近在学习PyTorch，想着光看不动手不行，就尝试着对着torchvision.models下的ResNet实现写了一遍，顺带将ResNet复习了一下，并将在研读源码过程中自己做的一点笔记和理解简单的分享一下。

torchvision.models.resnet下有ResNet18、ResNet34、ResNet50、ResNet101、ResNet152五种结构的ResNet实现。

主要的区别在于:

1. 18层和34层的ResNet用到的结构是"两个3*3的conv层堆叠"的BasicBlock，而50，101，152层的ResNet用到了Bottleneck结构，所以源码中也分别有BasicBlock和Bottleneck的实现。
2. 相同的BasicBlock/Bottleneck重复堆叠的次数不同。

贴一张ResNet的结构图可能会更清晰一些：

从而主体核心代码ResNet类的初始化形式即为：

[code]# block:BasicBlock or Bottleneck
# layers: 根据自己需要搭积木。ResNet50的设置为layers=[3, 4, 6, 3]
# num_classes:根据自己需要设置最终fc层输出的分类数目，默认是ImageNet的1000类
def __init__(self, block, layers, num_classes=1000):

所以接下来主要重点解读Bottleneck类和ResNet类的实现。

一、Bottleneck类源码解读

同样，先贴一张Bottleneck的结构图（图片来自网络，依据tensorflow画的，PyTorch的image Tensor应该是(N,C,H,W))，图和代码相结合更有助于理解。

[code]class Bottleneck(nn.Module):
    '''Bottleneck block'''

    def __init__(self, inplanes, planes, stride=1, downsample=None):
# inplanes和planes的意义可以结合1*1卷积核的作用来理解：
# 先通过1*1卷积核将输入的Tensor的channel数从inplanes降低到planes，压缩特征(conv1);
# 然后通过3*3卷积核抽取特征，channel数不变，仍是planes(conv2);
# 最后再次通过1*1卷积核将输出的Tensor的channel数从planes提升到4*planes(conv3);
        super(Bottleneck, self).__init__()
        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
        self.bn1 = nn.BatchNorm2d(planes)
        # 此处用3*3kernel作用时，若stride=2相当于下采样一次，此时在两个通道相加时，恒等映射通道需要做一次下采样，注意仅此处用到了stride参数
        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(planes)
        self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
        self.bn3 = nn.BatchNorm2d(planes * 4)
        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample
        self.stride = stride
    
    def forward(self, x):
        # 恒等映射块
        residual = x
        
        # 通过卷积层的正向传播
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu(out)

        out = self.conv3(out)
        out = self.bn3(out)
        
        # 判断是否需要下采样，即是否两个通道的output_channel是一致的
        if self.downsample is not None:
            residual = self.downsample(x)
        
        # 注意此处是两个并行通道相加后再做ReLU
        # 注意是相加，不是concat起来！
        out += residual
        out = self.relu(out)
        
        return out

二、ResNet类源码解读

[code]class MyResNet50(nn.Module):
'''My own ResNet'''
def __init__(self, bottleneck, layers, num_classes=1000):
super(MyResNet50, self).__init__()

self.inplanes = 64

# 第一部分
# 为了使图像224*224 -> 112*112,需要设置stride=2,padding=3
self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3)
self.bn1 = nn.BatchNorm2d(64)
self.relu = nn.ReLU(inplace=True)

# 第二部分
# 为了使图像112*112 -> 56*56,需要设置stride=2,padding=1
self.maxpool1 = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
# 第二部分图像size减半的工作由maxpooling层已完成，所以stride=1不再减半
self.layer1 = self._make_layer(bottleneck, 64, layers[0], stride=1)

# 第三部分
# stride=2 图像size减半
self.layer2 = self._make_layer(bottleneck, 128, layers[1], stride=2)

# 第四部分
# stride=2 图像size减半
self.layer3 = self._make_layer(bottleneck, 256, layers[2], stride=2)

# 第五部分
# stride=2 图像size减半
self.layer4 = self._make_layer(bottleneck, 512, layers[3], stride=2)

# 第六部分
# avgpooling 使 N*2048*7*7的Tensor -> N*2048*1*1的Tensor
self.avgpool = nn.AvgPool2d(7, stride=1)
self.fc = nn.Linear(2048, num_classes)

# 初始化各Module的参数
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))
elif isinstance(m, nn.BatchNorm2d):
m.weight.data.fill_(1)
m.bias.data.zero_()

def _make_layer(self, bottleneck, planes, blocks, stride=1):
downsample = None

if stride != 1 or self.inplanes != planes * 4:
# 注意设置downsample的条件
# 当满足stride != 1时，是需要改变恒等映射Tensor的h,w
# 而当满足self.inplanes != planes * 4时，时需要改变恒等映射Tensor的channel
# 当且仅当stride=1,self.inplanes=planes*4时，输入的Tensor的size和经过卷积层得到的Tensor的size才是一致的，不需要下采样
downsample = nn.Sequential(
nn.Conv2d(self.inplanes, planes * 4, kernel_size=1, stride=stride, bias=False),
nn.BatchNorm2d(planes * 4)
)

# 将bottleneck block堆叠成layer
layers = []
layers.append(bottleneck(self.inplanes, planes, stride, downsample))
# 此处是关键，下面这句起到了连接一个layer里重复的bottleneck的作用
# 将下一个bottleneck的input_channel值设为上一个bottleneck的output_channel值
# 所以planes是一个bottleneck层的input_channel数，
# self.inplanes=planes*4是该层每一个bottleneck块的output_channel数，也是该层的output_channel数
self.inplanes = planes * 4
for i in range(1, blocks):
# 每个重复Bottleneck layer仅第一个bottleneck需要改变图像分辨率以及恒等映射的channel值
layers.append(bottleneck(self.inplanes, planes))

return nn.Sequential(*layers)

def forward(self, x):
x = self.conv1(x)
x = self.bn1(x)
x = self.relu(x)
x = self.maxpool1(x)

x = self.layer1(x)
x = self.layer2(x)
x = self.layer3(x)
x = self.layer4(x)

x = self.avgpool(x)
# 在进行fc层前，由于此时Tensor的size为(N,C,1,1)，而fc要求输入为(N,input_size)
# 所以需要固定N，即x.size(0)，然后把后面的(C,1,1)拉成一个值
x = x.view(x.size(0), -1)
x = self.fc(x)

return x

如有疑问或者本人理解不正确的地方，希望大家能够在评论里提出，一起交流探讨，谢谢！