[논문구현]Pytorch로 ResNet 구현

이번 포스팅에서는 ResNet 을 Pytorch로 구현해보겠습니다.

이전 ResNet 논문 리뷰 링크
https://dreamrunning.tistory.com/10

 

 

ResNet 은 2015년에 나온 CNN 모델로서, 현재도 다양한 분야에서 많이 쓰이는 모델입니다. 특히 ResNet 에서 처음 제안하고 사용한 Residual Learning 개념은 현재에도 다양하게 응용되고 있습니다.

 

 

우선 필요 패키지를 불러옵니다.

import torch
import torch.nn as nn
import torch.nn.functional as F

 

1. Basic Building Block 및 Bottleneck Building Block

왼쪽: Basic Building Block, 오른쪽: Bottleneck Building Block

ResNet 의 가장 기본 구조인 Basic Building Block 및 Bottleneck Building Block 을 먼저 구현하겠습니다.

 

1.1 Basic Building Block

Basic Building Block 은 kernel size 가 3x3 인 convolution layer 두 개와 skip connection 으로 이루어진 Residual network 구조를 가집니다.

논문 리뷰에서 언급했듯이, 한 번의 convolution layer 를 거친 후 Normalization 을 하고, 이후 activation funtion 을 거칩니다. 따라서 한 개의 convolution layer 당 다음 구조를 가져야 합니다.

• 3x3 Convolution layer
• Normalization layer (Batch Normalization)
• Activation fuction (ReLU)

Batch Normalization layer 에서 bias 를 가지므로, Convolution layer 에서는 bias=False 로 설정합니다.

 

이를 코드로 구현한다면, Basic Building Block 은 다음과 같은 인스턴스 속성을 가집니다.

class BasicBlock(nn.Module):
    def __init__(self, ... ):
        super().__init__()
        ...

        self.conv = nn.Conv2d(in_ch, out_ch, kernel_size=3, stride=1, padding=1, bias=False)
        self.norm = nn.BatchNorm2d(out_channels)
        self.actv = nn.ReLU()

 

Resnet 에서는 feature map size 를 1/2 로 만들 시 stride=2 를 사용합니다. 따라서 feature map size 를 줄일 때 첫 번째 Convolution layer 의 stride 값을 2로 만들어야 합니다.

 

feature map size 를 절반으로 줄일 시

• convolutoin layer 1 : stride=2
• convolution layer 2 : stride=1 (고정)

이는 파라메터로 조정할 수 있으며, 코드로 구현 시 다음과 같이 구현할 수 있습니다.

class BasicBlock(nn.Module):
    def __init__(self, stride=2, ... ):
        super().__init__()
        ...
		
        # 첫 번째conv layer의 stride 는 파라메터에 의해 변경
        self.conv1 = nn.Conv2d(in_ch, out_ch, kernel_size=3, stride=stride, padding=1, bias=False)
        # 두 번째conv layer의 stride 는 1 로 고정
        self.conv2 = nn.Conv2d(out_ch, out_ch, kernel_size=3, stride=1, padding=1, bias=False)

 

 

Skiip Connection 의 경우, 기본적으로 Identitiy 이지만, 위 Layers 를 거친 후 출력 벡터의 차원이 바뀔 시 1x1 Convolution filter 를 이용하여 Skip Connection 출력과 벡터의 차원을 일치시킵니다.

• Skip Connection: Identity (입력과 출력 차원이 동일 시)
• Skip Connection: 1x1 Convolution filter (입력과 출력 차원이 다를 시)

이를 코드로 구현하면, 다음과 같습니다.

class BasicBlock(nn.Module):
    def __init__(self, ... ):
        super().__init__()
        ...

        self.residual = None
        # 입력과 출력의 차원이 불일치 시
        if in_channels != out_channels:
            self.residual = nn.Sequential(
                nn.Conv2d(in_ch, out_ch, kernel_size=1, stride=stride, padding=0, bias=False),
                nn.BatchNorm2d(out_channels)
            )
        # 입력과 출력의 차원이 동일할 시
        else:
            self.residual = nn.Identity()

 

이를 종합해서, 다음과 같이 Basic Building Block 를 구현합니다. expansion_factor 는 사용하지 않는 파라메터지만 ResNet Class 에서 호환성을 위해 적습니다.

class BasicBlock(nn.Module):
    def __init__(self, 
                 in_ch: int, 
                 out_ch: int, 
                 stride: int = 1,
                 padding: int = 1,
                 expansion_factor: int = 1
                 ):
        super().__init__()

        self.conv1 = nn.Conv2d(in_ch, out_ch, kernel_size=3, stride=stride, padding=padding, bias=False)
        self.conv2 = nn.Conv2d(out_ch, out_ch, kernel_size=3, stride=1, padding=padding, bias=False)
        self.norm1 = nn.BatchNorm2d(out_channels)
        self.norm2 = nn.BatchNorm2d(out_channels)
        self.actv = nn.ReLU()

        self.residual = None
        if in_channels != out_channels:
            self.residual = nn.Sequential(
                nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=stride, padding=0, bias=False),
                nn.BatchNorm2d(out_channels)
            )
        else:
            self.residual = nn.Identity()

    def forward(self, x):
        # skip connection
        residual = self.residual(x)

        x = self.conv1(x)
        x = self.norm1(x)
        x = self.actv(x)
        x = self.conv2(x)
        x = self.norm2(x)
        x = x + residual

        return self.actv(x)

 

1.2 Bottleneck block

Bottleneck Block 은 1x1 convolution layer 두개와 그 사이 3x3 convolution layer 로 구성되어 있습니다. Basic Block 과 마찬가지로 Normalization layer 가 사이사이마다 들어가 있는 형태입니다. 기초적인 형태는 다음과 같습니다.

• 1x1 or 3x3 Convolution layer
• Normalization layer (Batch Normalization)
• Activation function (ReLU)

여기서, Basic Block 과 차이점은 마지막 1x1 Convolution layer 를 이용해 차원을 줄였다가 다시 늘립니다(복원). pytorch 공식 코드에서는 차원을 4배로 늘였다가 줄입니다. 공식 코드에서는 width = int(planes * (base_width / 64.0)) * groups 으로 Bottleneck Block 내부에서 줄여진 차원을 따로 계산합니다만 이 포스팅에서는 생략하겠습니다.

 

각 Bottleneck Block 마다 channel input 및 output 은 다음과 같습니다.

• 1x1 conv layer: in_channel -> inner_dim
• 3x3 conv layer: inner_dim -> inner_dim
• 1x1 conv layer: inner_dim -> inner_dim * expansion_factor(default=4)

공식 코드에서는 Bottleneck Block 기준 dimension 이 inner_dim 에 의해 설정됩니다.

 

또한 feature map size 를 절반으로 줄일 시, 3x3 convolution layer 의 stride 값을 조정합니다. 이를 코드로 구현하면 다음과 같습니다.

class BottleneckBlock(nn.Module):
    def __init__(self,
                 ...
                 inner_dim: int
                 stride: int = 1,
                 expansion_factor: int = 4
                 ...
                 ):

        self.conv1 = nn.Conv2d(in_ch, inner_dim, kernel_size=1, stride=1, bias=False)
        self.conv2 = nn.Conv2d(inner_dim, inner_dim, kernel_size=3, stride=stride, padding=1)
        self.conv3 = nn.Conv2d(inner_dim, inner_dim * expansion_factor, kernel_size=1, stride=1, bias=False)

        self.norm1 = nn.BatchNorm2d(inner_dim)
        self.norm2 = nn.BatchNorm2d(inner_dim)
        self.norm3 = nn.BatchNorm2d(inner_dim * expansion_factor)
        ...

 

Skip Connection 의 경우 Basic Block 과 기본적으로 동일하게 구현되지만, block dimension 이 달라지는 경우가 stride=2 일때 및 in_channel != out_channel(inner_dim * expansion_factor) 두 가지 경우가 있으므로 추가적으로 수정해야 합니다.

 

Bottleneck Block 전체를 코드로 구현하면 다음과 같습니다.

class BottleneckBlock(nn.Module):
    def __init__(self,
                 in_ch: int,
                 inner_dim: int,
                 stride: int = 1,
                 padding: int = 1,
                 expansion_factor: int = 4
                 ):

        self.conv1 = nn.Conv2d(in_ch, inner_dim, kernel_size=1, stride=1, bias=False)
        self.conv2 = nn.Conv2d(inner_dim, inner_dim, kernel_size=3, stride=stride, padding=padding)
        self.conv3 = nn.Conv2d(inner_dim, inner_dim * expansion_factor, kernel_size=1, stride=1, bias=False)

        self.norm1 = nn.BatchNorm2d(inner_dim)
        self.norm2 = nn.BatchNorm2d(inner_dim)
        self.norm3 = nn.BatchNorm2d(inner_dim * expansion_factor)

        self.actv = nn.ReLU()

        self.residual = None
        if stride == 2 or in_ch != inner_dim * expansion_factor:
            self.residual = nn.Sequential(
                nn.Conv2d(in_ch, inner_dim * expansion_factor, kernel_size=1, stride=stride, padding=0),
                nn.BatchNorm2d(inner_dim * expansion_factor)
            )
        else:
            self.residual = nn.Identity()
        
    
    def forward(self, x):
        # skip connection
        residual = self.residual(x)
        # 1x1 convolution layer
        x = self.conv1(x)
        x = self.norm1(x)
        x = self.actv(x)
        # 3x3 convolution layer
        x = self.conv2(x)
        x = self.norm2(x)
        x = self.actv(x)
        # # 1x1 convolution layer
        x = self.conv3(x)
        x = self.norm3(x)
        x = x + residual

        return self.actv(x)

 

이렇게 Basic Block 과 Bottleneck Block 을 구성하였습니다.

 

 

2. ResNet Class

이제 ResNet Class 를 구현해보겠습니다. 논문에서 ResNet 은 5가지 종류를 언급하였습니다.

 

이 5가지 ResNet 을 만들 때 Basic Block 또는 Bottleneck Block 을 일일이 코드를 쳐서 구현해도 되지만, 따로 함수를 만들어서 구성되게 해보겠습니다.

 

우선 ResNet Class 내 필요한 인스턴스 속성은 입력차원, Block 종류, expansion_factor 입니다. 이 세가지가 필요한 이유는 다음과 같습니다.

• 입력차원: 각 Block의 input dimension 을 정의하기 위해서 필요, 처음 값 64, Block 이 반복될 때 마다 갱신
• Block: Basic Block 과 Bottleneck Block 중 어떤 Block을 사용할것인지 정의하기 위해 필요
• expansion_factor: Basic Block 일 때 값은 1, Bottleneck Block 일 때는 4

expansion_factor 는 Block 종류에 따라 값이 종속되므로 다음과 같이 코드를 구현할 수 있습니다.

class ResNet(nn.Module):
    def __init__(self, block, ...):
        super().__init__()

        self.in_channels = 64
        self.block = block
        self.expantion_factor = 4 if self.block == BottleneckBlock else 1
        ...

이 때 파라메터 block 은 BasicBlock class 또는 Bottleneck Class 를 직접 받습니다.

 

 

ResNet 은 수많은 Block 으로 레이어를 구성합니다. 위 표에 따라 생성되는 메소드를 make_layer 함수라 하고 이를 구현해보겠습니다.

 

 

2.1 make_layer method

make_layer 메소드는 위 표에서 Basic Block 및 Bottleneck Block 을 반복하여 ResNet 의 핵심 Convolution layer 를 구성하게 만드는 역할입니다.

 

핵심 구조는 for 문을 통하여 num_block 만큼 block 을 반복시켜 layers 에 append 하는 형식입니다. 이 때 고려해야할 점은, 첫번째 block 의 input dimension 은 self.in_channels 이고 그 이후 반복되는 block 의 input 은 inner_dim * expansion_factor 라는 것입니다.

• 1st Block : in_channels -> inner_dim * expansion_factor
• 2nd Block 이후: inner_dim * expansion_factor -> inner_dim * expansion_factor

expansion_factor 는 Basic Block 일 경우 1 이므로 입력 및 출력 차원에 영향을 주지 않습니다. Bottleneck Block 인 경우에만 영향을 줍니다.

 

코드로 구현하면 다음과 같습니다.

class ResNet(nn.Module):
    def __init__(self, block):
        super().__init__()
		
        self.in_channels = 64
        self.block = block
        self.expantion_factor = 4 if self.block == BottleneckBlock else 1
        ...

    def make_layer(self, channels = 64, num_block = 3, stride = 1):
        layers = nn.ModuleList([])
        # 1st Block
        layers.append(self.block(self.in_channels, channels, stride))
        # self.in_channels 갱신
        self.in_channels = channels * self.expantion_factor 
        # 2nd 이후 Block
        for _ in range(1, num_block):
            layers.append(self.block(self.in_channels, channels))
    
        return nn.Sequential(*layers)

 

make_layer 메소드를 사용하여 레이어를 선언하면 다음과 같습니다. 이 떄 각 레이어의 내부 차원은 모든 ResNet 마다 동일하므로 레이어를 선언 시 필요 숫자를 파라메터로 집어넣습니다.

class ResNet(nn.Module):
    def __init__(self, block, ...):
        super().__init__()
        ...
        
        self.layer1 = self.make_layer(channels=64, num_block=3)

 

또한 Block이 반복될 수를 정해야 하므로, 초기화 함수에서 리스트 형태로 파라메터를 받도록 합니다. 이를 종합하면 다음과 같습니다.

class ResNet(nn.Module):
    def __init__(self, block, layers_num = [3, 4, 6, 3], ...):
        super().__init__()

        self.in_channels = 64
        self.block = block
        self.expantion_factor = 4 if self.block == BottleneckBlock else 1

        # 레이어 선언
        self.layer1 = self.make_layer(channels=64, num_block=layers_num[0])
        self.layer2 = self.make_layer(channels=128, num_block=layers_num[1], stride=2) # 여기서부터 downsampling
        self.layer3 = self.make_layer(channels=256, num_block=layers_num[2], stride=2)
        self.layer4 = self.make_layer(channels=512, num_block=layers_num[3], stride=2)
        ...

    # Block 반복 레이어 생성 메소드
    def make_layer(self, channels = 64, num_block = 3, stride = 1):
        layers = nn.ModuleList([])
        # 1st Block
        layers.append(self.block(self.in_channels, channels, stride))
        # self.in_channels 갱신
        self.in_channels = channels * self.expantion_factor 
        # 2nd 이후 Block
        for _ in range(1, num_block):
            layers.append(self.block(self.in_channels, channels))
    
        return nn.Sequential(*layers)

 

 

2.2 ResNet Class 완성시키기

make_layer 메소드를 사용하여 Block 을 반복시켜 layer 를 생성하도록 코드를 만들어봤습니다. 이제 ResNet 에 필요한 나머지 부분을 완성시켜보겠습니다.

 

ResNet은 처음 이미지를 입력으로 받을 시 7x7 Convolution layer 를 거치고 이후 MaxPool layer 를 거칩니다.

• Convolution Layer: kernel size=7x7, in_ch=3, out_ch=self.in_channels, stride=2, padding=3, bias=False
• Maxpool Layer: kernel size=3x3, stride=2, padding=1

 

make_layer 로 만들어진 레이어들은 거친 후에는 AveragePooling 을 한 뒤 Linear Layer 를 사용하여 이미지 클래스를 분류합니다. 이 사이에 torch.flatten 으로 feature_map 을 1차원으로 만듭니다.

• Avgpool: kernel size=1x1
• Fully-connected Layer: in_dim=512*expansion_factor, out_dim=1000

이를 종합하여 ResNet 을 완성시키면 다음과 같습니다.

class ResNet(nn.Module):
    def __init__(self, block, layers_num = [3, 4, 6, 3]):
        super().__init__()

        self.in_channels = 64
        self.block = block
        self.expantion_factor = 4 if self.block == BottleneckBlock else 1

        # conv $ maxpool
        self.conv = nn.Conv2d(3, self.in_channels, kernel_size=7, stride=2, padding=3, bias=False)
        self.norm = nn.BatchNorm2d(self.in_channels)
        self.actv = nn.ReLU()
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)

        # 레이어 선언
        self.layer1 = self.make_layer(channels=64, num_block=layers_num[0])
        self.layer2 = self.make_layer(channels=128, num_block=layers_num[1], stride=2) # 여기서부터 downsampling
        self.layer3 = self.make_layer(channels=256, num_block=layers_num[2], stride=2)
        self.layer4 = self.make_layer(channels=512, num_block=layers_num[3], stride=2)

        # avgpool and fc layer
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(512*self.expantion_factor, 1000)

    # Block 반복 레이어 생성 메소드
    def make_layer(self, channels = 64, num_block = 3, stride = 1):
        layers = nn.ModuleList([])
        # 1st Block
        layers.append(self.block(self.in_channels, channels, stride))
        # self.in_channels 갱신
        self.in_channels = channels * self.expantion_factor 
        # 2nd 이후 Block
        for _ in range(1, num_block):
            layers.append(self.block(self.in_channels, channels))
    
        return nn.Sequential(*layers)
    
    def forward(self, x):
        x = self.conv(x)
        x = self.norm(x)
        x = self.actv(x)
        x = self.maxpool(x)

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

        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        x = self.fc(x)

        return x

 

이렇게 ResNet Class 를 완성시킬 수 있습니다.