[pytorch] torch.nn.Module Build

오늘은 torch.nn.Module 패키지를 활용한 Deep Neural Network를 구축해보고, 내가 얻은 인사이트를 공유하겠습니다.

 

1. 저는 torch.nn.Module 패키지로 보통 클래스를 생성하는데, nn.Module을 기반클래스로 상속하여 파생클래스인 Model(CNN)을 생성하여 구축합니다. (여기선 CNN이라는 Class 명을 사용했습니다.)

 

import torch
import torch.nn as nn

class CNN(nn.Module):
    def __init__(self, img_size, num_class):
        super(CNN, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(3, 32, 3, 1, 1),
            nn.LeakyReLU(0.2),
            nn.Conv2d(32, 64, 3, 1, 1),
            nn.LeakyReLU(0.2),
            nn.MaxPool2d(2, 2),
            nn.Conv2d(64, 128, 3, 1, 1),
            nn.LeakyReLU(0.2),
            nn.Conv2d(128, 256, 3, 1, 1),
            nn.LeakyReLU(0.2),
            nn.MaxPool2d(2, 2),
            nn.Conv2d(256, 512, 3, 1, 1),
            nn.LeakyReLU(0.2),
            nn.MaxPool2d(2,2),
            nn.Conv2d(512, num_class, 3, 1, 1),
            nn.LeakyReLU(0.2)
        )
        
        self.avg_pool = nn.AvgPool2d(img_size // 8)
        self.classifier = nn.Linear(num_class, num_class)

    def forward(self, x):
        features = self.conv(x)
        flatten = self.avg_pool(features).view(features.size(0), -1)
        output = self.classifier(flatten)
        return output, features

 

상기 코드를 설명하자면:

1. self.conv 변수를 nn.Sequential() Class의 생성자로 선언함.

 

2. nn.Sequential() 내부적으로 argment를 nn.Conv2d, nn.LeakyReLU, nn.MaxPool2d를 이용하여 간단한 DNN Layer 구조를 생성

 

3. 그 다음 종단의 Average Pooling과 Fully Conneted Layer를 self.avg_pool = nn.AvgPool2d(kernelsize), self.classifier = nn.Liner(input channel, output channel)로 생성한다.

 

4. 마지막으로 def forward(self, x): 함수를 선언하여 실제로 test data를 받아왔을때 conv -> avg_pool -> fc 순으로 투입되게 layer 순서를 정의한다.

 

* 여기서 핵심은 아래와 같이 test tensor를 전달했을때 자동으로 forward 함수로 전달되어 구성된 Network를 통과하면서 학습이 진행된다.

 

cnn = model.CNN(img_size=img_size, num_class=numclass)
logit, last_feature = cnn(test_tensor)

 

즉, x에 test_tensor가 부여됨.

 

    def forward(self, x):
        features = self.conv(x)
        flatten = self.avg_pool(features).view(features.size(0), -1)
        output = self.classifier(flatten)

 

디버그해보시면 이런 순서로 model이 진행되는것을 알 수 있습니다.

 

감사합니다.