resnet2d.py

Go to the documentation of this file.

import torch
import torch.nn as nn
 
#ResNet2d의 베이스블럭
class BasicBlock(nn.Module):
    def __init__(
        self,
        inplanes: int,
        planes: int,
        stride: int = 1,
        downsample = None
    ) -> None:
        super(BasicBlock, self).__init__()
        
        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=(3, 1), stride=stride, padding=(1, 0), bias=False)
        self.bn1 = nn.BatchNorm2d(planes)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = nn.Conv2d(inplanes, planes, kernel_size=(3, 1), stride=1, padding=(1, 0), bias=False)
        self.bn2 = nn.BatchNorm2d(planes)
        self.downsample = downsample
        self.stride = stride
 
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        
        identity = x
        
        if self.downsample is not None:
            identity = self.downsample(x)
 
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)
 
        out = self.conv2(out)
        out = self.bn2(out)
 
        out += identity
        out = self.relu(out)
        
        return out
 
class ResNet2d(nn.Module):
    """
    이미지 특징을 사용하기 위한 모델
    입력을 T x 1 로 가정하고 만듦
    """
    def __init__(
        self,
        in_channels, 
        stage_blocks,
        input_key,
    ) -> None:
        """
        args:
            in_channels (int) : 입력 채널수
            stage_blocks (List) : 스테이지 마다 쌓는 블럭의 수를 설정
            input_key (str) : 모듈의 inference에서 사용하는 입력데이터의 키값
 
        """
        super(ResNet2d, self).__init__()
 
        block = BasicBlock
 
        self.inplanes = in_channels 
        self.input_key = input_key
        
        # input block
        self.conv1 = nn.Conv2d(in_channels, self.inplanes, kernel_size=1, stride=1, padding=0, bias=True)
        self.bn1 = nn.BatchNorm2d(self.inplanes)
        self.relu = nn.ReLU(inplace=True)
        
        # residual blocks
        self.layer1 = self._make_layer(block, in_channels, stage_blocks[0], stride=1)
        self.layer2 = self._make_layer(block, in_channels, stage_blocks[1], stride=2)
        self.layer3 = self._make_layer(block, in_channels, stage_blocks[2], stride=2)
        self.layer4 = self._make_layer(block, in_channels, stage_blocks[3], stride=1)
        
 
    def _make_layer(self, block, planes: int, blocks: int, stride: int = 1) -> nn.Sequential:
        
        downsample = None
        
        # downsampling 필요할경우 downsample layer 생성
        if stride != 1 or self.inplanes != planes:
            downsample = nn.Sequential(
                nn.Conv2d(self.inplanes, planes, kernel_size=1, stride=stride),
                nn.BatchNorm2d(planes)
            )
 
        layers = []
        layers.append(block(self.inplanes, planes, stride, downsample))
        self.inplanes = planes
        for _ in range(1, blocks):
            layers.append(block(self.inplanes, planes))
 
        return nn.Sequential(*layers)
 
    def forward(self, sample):
        """
        args:
            sample (dict)) : 입력 데이터, self.input_key에 해당하는 키가 있어야함
                self.input_key의 아이템은 Tensor 타입 -> shape (B, C, T, 1)
                    B : 배치 크기
                    C : 입력 채널
                    T : 시간
 
        return (Tensor):
            특정 해상도의 특징 벡터 -> shape (B, C_o, T_o, 1)
                B : 배치 크기
                C_o : 채널
                T_o : 시간
        """
        x = sample[self.input_key]
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
 
        return x