본문 바로가기
Machine Learning/개념 정리

[파이토치] 4. 여러 모듈 사용

by W_log 2023. 12. 29.

3번 과정까지의 글은 신경망 구조 자체가 간단한 모델이라고 볼 수 있다. 하지만 실제 논문들의 코드를 보다보면, 생각보다 복잡한 형태로 구현되는 경우가 있다. 이를 효율적으로 구성하기 위해 클래스를 여러개 만드는 방법을 소개한다.

 

또한, 사실 신경망 구조 자체는 연구자가 한땀한땀 설정하는 것이기 때문에 한번에 하나의 모델이 아니라 여러 개의 모델을 돌려놓기도 하는데 이에 대해서도 코드로 살펴보려고 한다.

 

1.  신경망 구조 모듈화하기

 

출처 : https://arxiv.org/abs/1409.4842

 

 

구글에서 만든 이미지 분류 모델 중 하나인 GoogLeNET이 있는데 구글넷에는 위의 그림처럼 Inception module이 여러개 붙어있는 구조이다. 하지만 class GoogLeNET의 __init__ 에 다 넣기보다는 Inception module을 별도 클래스로 지정하고, 이 모듈을 사용해서 붙이는 형태로 구성하는 형태로 코드를 짤 수 있다.

 

 

class Inception(nn.Module):
    def __init__(self, in_channels, kernel_1_x, kernel_3_in, kernel_3_x, kernel_5_in, kernel_5_x, pool_channels):
        super(Inception, self).__init__()
        # 1x1 conv
        self.branch1 = nn.Sequential(
            nn.Conv2d(in_channels, kernel_1_x, kernel_size=1), # 1x1 conv
            nn.BatchNorm2d(kernel_1_x),
            nn.ReLU(True),
        )

        # 1x1 conv -> 3x3 conv branch
        self.branch2 = nn.Sequential(
            nn.Conv2d(in_channels, kernel_3_in, kernel_size=1), # 1x1 conv
            nn.BatchNorm2d(kernel_3_in),
            nn.ReLU(inplace = True),
            nn.Conv2d(kernel_3_in, kernel_3_x, kernel_size=3, padding=1), # 3x3 conv
            nn.BatchNorm2d(kernel_3_x),
            nn.ReLU(inplace = True),
        )

        # 1x1 conv -> 5x5 conv branch
        self.branch3 = nn.Sequential(
            nn.Conv2d(in_channels, kernel_5_in, kernel_size=1), # 1x1 conv
            nn.BatchNorm2d(kernel_5_in),
            nn.ReLU(True),
            nn.Conv2d(kernel_5_in, kernel_5_x, kernel_size=5, padding=2), # 5x5 conv
            nn.BatchNorm2d(kernel_5_x),
            nn.ReLU(True),
        )

        # 3x3 pool -> 1x1 conv branch
        self.branch4 = nn.Sequential(
            nn.MaxPool2d(3, stride=1, padding=1), # 3x3 max pooling
            nn.Conv2d(in_channels, pool_channels, kernel_size=1), # 1x1 conv
            nn.BatchNorm2d(pool_channels),
            nn.ReLU(True),
        )

    def forward(self, x):
        '''
        INPUT :
            x = [batch_size, channel, height, width]
        OUTPUT :
            output = [batch_size, num_classes]
        '''
        y1 = self.branch1(x) # [batch_size, kernel_1_x, height, width]
        y2 = self.branch2(x) # [batch_size, kernel_3_x, height, width]
        y3 = self.branch3(x) # [batch_size, kernel_5_x, height, width]
        y4 = self.branch4(x) # [batch_size, pool_channels, height, width]
        # 병렬로 계산된 convolution 들을 모두 합쳐줍니다.
        return torch.cat([y1,y2,y3,y4], dim = 1)  # [batch_size, kernel_1_x + kerenl_3_x + kernel_5_x + pool_channels, height, width]


class GoogLeNet(nn.Module): # googlenet 논문의 구조를 그대로 반영하였습니다.
    def __init__(self, num_classes, aux_logits = True):
        super(GoogLeNet, self).__init__()

        self.num_classes = num_classes
        self.aux_logits = aux_logits
        #Inception layer로 들어가기 전에 나오는 Layer
        self.pre_layers = nn.Sequential(
            nn.Conv2d(3, 192, kernel_size=3, padding=1),
            nn.BatchNorm2d(192),
            nn.ReLU(True),
        )
        #in_channels, kernel_1_x, kernel_3_in, kernel_3_x, kernel_5_in, kernel_5_x, pool_channels
        self.inception_1 = Inception(192,  64,  96, 128, 16, 32, 32) # inception 1, inception 2 이런 식으로 변수 이름 변경
        self.inception_2 = Inception(256, 128, 128, 192, 32, 96, 64)

        self.max_pool = nn.MaxPool2d(3, stride=2, padding=1)

        self.inception_3 = Inception(480, 192,  96, 208, 16,  48,  64)  #kernel_1_x + kernel_3_x + kernel_5_x + pool_channels로 다음게 넘어감
        self.inception_4 = Inception(512, 160, 112, 224, 24,  64,  64)
        self.inception_5 = Inception(512, 128, 128, 256, 24,  64,  64)
        self.inception_6 = Inception(512, 112, 144, 288, 32,  64,  64)
        self.inception_7 = Inception(528, 256, 160, 320, 32, 128, 128)

        self.inception_8 = Inception(832, 256, 160, 320, 32, 128, 128)
        self.inception_9 = Inception(832, 384, 192, 384, 48, 128, 128)


        self.avgpool = nn.AvgPool2d(8, stride=1)
        self.linear = nn.Linear(1024, self.num_classes)
        self.softmax = nn.LogSoftmax(dim = 1)

    def forward(self, x):
        '''
        INPUT :
            x = [batch_size, channel, height, width]
        OUTPUT :
            output = [batch_size, num_classes]
        '''
        x = self.pre_layers(x) # [batch_size, 192, height, width]
        x = self.inception_1(x) # [batch_size, 256, height, width]
        x = self.inception_2(x) # [batch_size, 480, height, width]
        x = self.max_pool(x) # [batch_size, 480, height//2, width//2]
        x = self.inception_3(x) # [batch_size, 512, height//2, width//2]
        x = self.inception_4(x) # [batch_size, 512, height//2, width//2]
        x = self.inception_5(x) # [batch_size, 512, height//2, width//2]
        x = self.inception_6(x) # [batch_size, 528, height//2, width//2]
        x = self.inception_7(x) # [batch_size, 832, height//2, width//2]
        x = self.max_pool(x) # [batch_size, 832, height//4, width//4]
        x = self.inception_8(x) # [batch_size, 832, height//4, width//4]
        x = self.inception_9(x) # [batch_size, 832, height//4, width//4]
        x = self.avgpool(x) # [batch_size, 1024, height//4, width//4]
        x = x.view(x.size(0), -1) # [batch_size, 1024, 1, 1]
        x = self.linear(x) # [batch_size, num_classes]
        x = self.softmax(x) # [batch_size, num_classes]
        return x

    def count_parameters(self):
        return sum(p.numel() for p in self.parameters() if p.requires_grad)

 

 

2.  다양한 신경망 구조를 학습시키기

 

이번에 활용할 모델은 VGG 모델인데, VGG는 layer 갯수를 늘려나가면서 학습시킬 수 있다. 여러 방법이 있지만, 미리 층의 갯수를 정의해놓고 그걸 인풋으로 받아서 처리해볼 수 있다.

 

우선 이렇게 미리 채널 수와 풀링을 정의한다. 여기서 11은 convolution 갯수 + fclayer 3개로 계산할 수 있다.

 

cfgs = {'VGG11' : [64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
    'VGG13' : [64, 64, 'M', 128, 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
    'VGG16' : [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M'],
    'VGG19' : [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512, 'M']}

 

 

이후에 Class를 만드는 과정에서 cfgs와 model_name을 인자로 받아온다. 우리가 정의한 cfgs는 dictionary로 Key값에는 모델명, values에는 모델 구조가 순서대로 구성되어 있다.

 

init에서 create_conv_layers라는 함수를 호출하고 이 함수에서는 모델 구조 리스트를 for문을 통해서 하나씩 생성하는 형태로 리스트에 맞는 모델을 만든다.

 

class VGG(nn.Module):
    def __init__(self, model_name, cfgs, in_channels=3, num_classes=10):
        #생략...
        self.cfgs = cfgs

        # cfgs에 저장된 모델 이름별 convolution layer의 output channel을 불러옵니다.
        self.conv_layers = self.create_conv_layers(cfgs[model_name]) # model은 생성하고자 하는 모델의 이름입니다. (type: str)

		#생략...


    def create_conv_layers(self, architecture):
        layers = []
        in_channels = self.in_channels # 1

        for x in architecture:
            if type(x) == int: # convolution layer
                out_channels = x

                layers += [nn.Conv2d(in_channels=in_channels, out_channels=out_channels,
                                     kernel_size=3, stride=1, padding=1),
                           nn.BatchNorm2d(x),
                           nn.ReLU(inplace = True)] # inplace = True는 원본 데이터에 덮어 씌우는 역할을 합니다.
                in_channels = x
            elif x == 'M': # max pooling
                layers += [nn.MaxPool2d(kernel_size=2)]
        # *는 가변인자로 layers 라는 변수 안의 값들을 모두 가져오는 역할을 합니다. 참고 : https://mingrammer.com/understanding-the-asterisk-of-python/
        return nn.Sequential(*layers)

 

이제 아래와 같이 정의해서 모델을 순차적으로 생성하고 훈련 코드도 유사하게 돌리면 작동한다.

 

for name in cfgs.keys():
    model_name = name
    model = VGG(model_name = model_name, cfgs = cfgs, in_channels = in_channels, num_classes = num_classes)
    print(f'{name} model parameters : ', model.count_parameters())
    
    

num_epochs = 100
in_channels = 3
num_classes = 10
lr = 1e-3
for name in cfgs.keys():
    model_name = name
    model = VGG(model_name = model_name, cfgs = cfgs, in_channels = in_channels, num_classes = num_classes).to(device)
    criterion = nn.NLLLoss()
    optimizer = optim.Adam(model.parameters(), lr = lr)
    model, valid_max_accuracy = training_loop(model, train_dataloader, valid_dataloader, train_dataset, valid_dataset, criterion, optimizer, device, num_epochs, patience, model_name)
    print(f'{model_name} valid max accuracy : ', valid_max_accuracy)
    test(model, model_name, test_dataloader, device)
    print('---'*20)
    print()