💡

[Pytorch] ResNetの実装

2022/11/06に公開

はじめに

pytorch で ResNet50, ResNet101, ResNet152 の実装をしてみました!

import torch
import torch.nn as nn

残差ブロックの実装

#残差ブロックを作成するクラス
class Resblock(nn.Module):
  def __init__(self, in_channels, out_channels, identity_downsample=None, stride=1):
    super(Resblock, self).__init__()

    self.expansion = 4
    
    self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
    self.bn1 = nn.BatchNorm2d(out_channels)
    self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=stride, padding=1)
    self.bn2 = nn.BatchNorm2d(out_channels)
    self.conv3 = nn.Conv2d(out_channels, out_channels*self.expansion, kernel_size=1, stride=1, padding=0)
    self.bn3 = nn.BatchNorm2d(out_channels*self.expansion)

    self.identity_downsample = identity_downsample
    self.relu = nn.ReLU()
  
  def forward(self, x):
    identity = x.clone()

    x = self.conv1(x)
    x = self.bn1(x)
    x = self.relu(x)
    x = self.conv2(x)
    x = self.bn2(x)
    x = self.relu(x)
    x = self.conv3(x)
    x = self.bn3(x)

    if self.identity_downsample is not None:
      identity = self.identity_downsample(identity)
    x = x + identity
    
    x = self.relu(x)

    return x

ResNet 全体の実装

class ResNet(nn.Module):
  def __init__(self, block, layers, num_classes):
    super(ResNet, self).__init__()
    
    self.in_channels = 64
    self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2,padding=3)
    self.bn1 = nn.BatchNorm2d(64)
    self.relu = nn.ReLU()
    self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2)

    self.conv2_x = self._make_layer(block, layers[0], first_conv_out_channels=64, stride=1)
    self.conv3_x = self._make_layer(block, layers[1], first_conv_out_channels=128, stride=2)
    self.conv4_x = self._make_layer(block, layers[2], first_conv_out_channels=256,stride=2)
    self.conv5_x = self._make_layer(block, layers[3], first_conv_out_channels=512, stride=2)

    self.avgpool = nn.AdaptiveAvgPool2d((1,1))
    self.fc = nn.Linear(512*4, num_classes)
  

  def forward(self, x):

    x = self.conv1(x)
    x = self.bn1(x)
    x = self.relu(x)
    x = self.maxpool(x)

    x = self.conv2_x(x)
    x = self.conv3_x(x)
    x = self.conv4_x(x)
    x = self.conv5_x(x)
    x = self.avgpool(x)
    x = x.reshape(x.shape[0], -1)
    x = self.fc(x)

    return x

    # 例 conv2_x のチャネル数の変化
    # in : 64
    # block(in_channels=64, first_conv_out_channels=64) out : 256
    # block(in_channels=256, first_conv_out_channels=64) out : 256
    # block(in_channels=256, first_conv_out_channels=64) out : 256
    # out : 256

  def _make_layer(self, block, num_resblocks, first_conv_out_channels, stride):
    identity_downsample = None
    layers = []

    # size 変更したいあるいは
    # チャネル数の変更が起こる場合残差接続でサイズとチャネル数が一致しなくなってしまうため
    # identity に conv層 を追加する。
    if stride != 1 or self.in_channels != first_conv_out_channels*4:
      identity_downsample = nn.Conv2d(self.in_channels, first_conv_out_channels*4, kernel_size=1, stride=stride)
    layers.append(block(self.in_channels, first_conv_out_channels, identity_downsample, stride))

    # 二つ目以降の残差ブロックでは常に、残差ブロックの入力チャネル数 in_channels が
    # その最初のconv層の出力チャネル数 first_conv_out_channels の4倍になっていることに注意
    self.in_channels = first_conv_out_channels*4

    for i in range(num_resblocks - 1):
      layers.append(block(self.in_channels, first_conv_out_channels))
    
    return nn.Sequential(*layers)

def get_ResNet50(block, num_classes):
  return ResNet(block, [3,4,6,3], num_classes)

def get_ResNet101(block, num_classes):
  return ResNet(block, [3,4,23,3], num_classes)

def get_ResNet152(block, num_classes):
  return ResNet(block, [3,8,36,3], num_classes)

forward だけ確認

10サンプルだけランダムで用意しました。

def main():
  net = get_ResNet152(Resblock, 1000)
  y = net(torch.randn(10, 3, 224, 224))
  print(y.size())
main()

結果は...大丈夫そうです!

torch.Size([10, 1000])

参考

https://arxiv.org/pdf/1512.03385.pdf
https://qiita.com/tchih11/items/377cbf9162e78a639958

Discussion