컴퓨터비전/pytorch

[Pytorch]ResNet(전차학습모델)

코딩의 숲 2023. 1. 13. 23:09

 

CNN은 레이어가깊어지면 기울기 소실이 쉽게 일어난다 그래서 전에있던값을 재탕함으로써 기울기를 살리는 기법

 

Block부분


class BasicBlock(nn.Module):
  def __init__(self,input_dim,output_dim, stride = 1):
    super(BasicBlock,self).__init__()
    self.conv1=nn.Conv2d(input_dim,output_dim,3,1,1)
    self.conv2=nn.Conv2d(output_dim,output_dim,3,1,1)
    self.bn1=nn.BatchNorm2d(output_dim)
    self.bn2=nn.BatchNorm2d(output_dim)


    self.relu1=nn.ReLU()
    self.relu2=nn.ReLU()
    self.shortcut=nn.Sequential()
    if input_dim!=output_dim or  stride != 1:
      self.shortcut=nn.Sequential(
          nn.Conv2d(input_dim,output_dim,3,1,1),
          nn.BatchNorm2d(output_dim)
      )
    
  def forward(self,x):
    out=self.conv1(x)
    out=self.bn1(out)
    out=self.relu1(out)
    out=self.bn2(out)
    out+=self.shortcut(x)

    out=self.relu2(out)

    
    return(out)

Model부분

class ResNet(nn.Module):
    def __init__(self, num_classes = 2):
        super(ResNet, self).__init__()
        self.in_planes = 16
        
        self.conv1 = nn.Conv2d(3, 16, kernel_size = 3, stride = 1, padding = 1, bias = False)
        self.bn1 = nn.BatchNorm2d(16)
        self.layer1 = BasicBlock(16,16,1)
        self.layer2 = BasicBlock(16,16,1)
        self.layer3 = BasicBlock(16,32,1)
        self.layer4 = BasicBlock(32,32,1)
        self.layer5 = BasicBlock(32,64,1)
        self.layer6 = BasicBlock(64,64,1)
        self.linear = nn.Linear(64*4*4, num_classes)
        
    
    
    def forward(self, x):
        out = F.relu(self.bn1(self.conv1(x)))
        out = self.layer1(out)

        out = self.layer2(out)

        out = self.layer3(out)
        out = self.layer4(out)
        out = self.layer5(out)
        out = self.layer6(out)

        out = F.avg_pool2d(out, 8)
        out = out.view(out.size(0), -1)
        out = self.linear(out)
        return out

AllexNet 논문을 읽고 구현해 봤는데 레이어를 깊게 쌓을수록 학습이 잘 안되었고 다른 논문을 뒤져본 결과 이 ResNet을 찾았다.