티스토리 뷰

AI

[논문 구현] ResNet 직접 구현해보기!

꿈더꿔 2022. 7. 27. 00:20

ResNet101을 기준으로 구현했다. bottleneck 을 class로 따로 떼었으며, first 변수로 사이즈를 줄이는 layer 여부를 구분했다. bottleneck은 conv 1x1 - batchnorm - relu - conv 3x3 - batchnorm - relu - conv 1x1 - skip connection - batchnorm - relu로 구성했다.

 

1. Setup

import torch
import torch.nn as nn
import numpy as np
import torchvision
import torchvision.transforms as transforms

from torch.utils.data import DataLoader
import torch.optim as optim
import time

import random
import torch.backends.cudnn as cudnn

seed = 2022
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
np.random.seed(seed)
cudnn.benchmark = False
cudnn.deterministic = True
random.seed(seed)

2. bottleneck

class bottleneck(nn.Module):
    def __init__(self, in_channels, mid_channels, out_channels, first = False, stride = 1):
        super(bottleneck, self).__init__()
        self.in_channels = in_channels
        self.mid_channels = mid_channels
        self.out_channels = out_channels
        self.first = first
        self.stride = stride

        self.conv1 = nn.Conv2d(
            in_channels = self.in_channels,
            out_channels = self.mid_channels,
            kernel_size=1,
        )
        self.batch1 = nn.BatchNorm2d(num_features=mid_channels)
        self.act1 = nn.ReLU()
        self.conv2 = nn.Conv2d(
            in_channels = self.mid_channels,
            out_channels = self.mid_channels,
            kernel_size = 3,
            stride = self.stride,
            padding = 1
        )
        self.batch2 = nn.BatchNorm2d(num_features=mid_channels)
        self.act2 = nn.ReLU()
        self.conv3 = nn.Conv2d(
            in_channels=self.mid_channels,
            out_channels=self.out_channels,
            kernel_size=1
        )
        self.act3 = nn.ReLU()
        self.batch3 = nn.BatchNorm2d(num_features=out_channels)
        self.proj = nn.Sequential(nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=1, stride = self.stride), nn.BatchNorm2d(num_features=out_channels))
        self.net = nn.Sequential(self.conv1, self.batch1, self.act1, self.conv2, self.batch2, self.act2, self.conv3)

    def forward(self, x):
        x2 = self.net(x)
        x2 = self.batch3(x2)
        if self.first:
            x = self.proj(x)
        x = torch.add(x, x2)
        x = self.act3(x)
        return x

한 bottleneck의 3x3 conv out channels는 이전 1x1 conv의 out channels와 같지만 in channels가 다르므로 mid channels 변수를 추가하였다.

first 변수를 통해 각 layer의 첫 bottleneck이면 True값을 주어 stride를 2로 바꿔 적용하여 layer마다 사이즈가 줄어들 수 있도록 했다.

First 가 아니면 Projection이 필요없으므로 그대로 identity 처럼 더하는 형태로 구현했다.

3. ResNet

class ResNet101(nn.Module):
    def __init__(self, in_channels = 3, size = 224, classes = 10):
        super(ResNet101, self).__init__()
        self.classes = classes
        self.in_channels = in_channels
        self.size = size

        self.conv1 = nn.Conv2d(in_channels=self.in_channels, out_channels=64, stride=2, kernel_size=7, padding=0)
        self.pool1 = nn.MaxPool2d(kernel_size=3, stride=2)

        self.conv2 = nn.Sequential(
            bottleneck(in_channels=64, mid_channels=64, out_channels=256, first=True),
            bottleneck(in_channels=256, mid_channels=64, out_channels=256),
            bottleneck(in_channels=256, mid_channels=64, out_channels=256),
        )

        self.conv3 = nn.Sequential(
            bottleneck(in_channels=256, mid_channels=128, out_channels=512, first=True, stride=2),
            bottleneck(in_channels=512, mid_channels=128, out_channels=512),
            bottleneck(in_channels=512, mid_channels=128, out_channels=512),
            bottleneck(in_channels=512, mid_channels=128, out_channels=512),
        )
        conv4 = [bottleneck(in_channels=512, mid_channels=256, out_channels=1024, first=True, stride=2)]
        conv4_2 = [bottleneck(in_channels=1024, mid_channels=256, out_channels=1024) for _ in range(22)]
        conv4.extend(conv4_2)
        self.conv4 = nn.Sequential(*conv4)

        self.conv5 = nn.Sequential(
            bottleneck(in_channels=1024, mid_channels=512, out_channels=2048, first=True, stride=2),
            bottleneck(in_channels=2048, mid_channels=512, out_channels=2048),
            bottleneck(in_channels=2048, mid_channels=512, out_channels=2048),
        )

        self.pool2 = nn.AdaptiveAvgPool2d((1, 1))
        self.features = nn.Sequential(
            self.conv1,
            self.pool1,
            self.conv2,
            self.conv3,
            self.conv4,
            self.conv5,
            self.pool2
        )
        self.flat = nn.Flatten()
        self.Linear = nn.Linear(in_features=2048, out_features=self.classes)
        self.soft = nn.Softmax(dim = 1)
        self.classifiers = nn.Sequential(
            self.Linear,
            self.soft
        )

    def forward(self, x):
        x = self.features(x)
        x = self.flat(x)
        x = self.classifiers(x)

        return x

conv4는 논문에서 23개의 bottleneck으로 구성되어있어서 list comprehension으로 구현했다.

 

CIFAR-10 을 통한 테스트코드는 아래 github에 업로드 해두었다.

https://github.com/kkt4828/reviewpaper/blob/630198644e880b635b31ad14875cf3f146df8081/ResNet/ResNet.py

 

GitHub - kkt4828/reviewpaper: 논문구현

논문구현. Contribute to kkt4828/reviewpaper development by creating an account on GitHub.

github.com