Pytorch CNN-LSTM 모델 설계

CNN-LSTM은 CNN으로 공간적 정보를 추출하고 LSTM으로 시계열 정보를 추출하여 동영상을 입력하는 테스크(행동인식 등)에서 좋은 성능을 얻을 수 있습니다.

그러나 I3D가 나오면서 CNN-LSTM모델의 단점을 거론하였는데요 바로 CNN에서 추출된 특징이 하이레벨 정보만 포함하고 있고 로우레벨은 LSTM으로 전달이 되지않아 더 높은 정확도를 얻을 수 없다는 점입니다.

우리는 I3D는 제쳐두고 CNN-LSTM 모델을 파이토치로 구현하는 방법을 알아봅시다.

 

import torch
import torch.nn as nn
from torch.autograd import Variable
from MobileNetV2 import MobileNetV2
from efficientnet_pytorch import EfficientNet

class EventDetector(nn.Module):
    def __init__(self, lstm_layers, lstm_hidden, bidirectional=True, dropout=True):
        super(EventDetector, self).__init__()
        self.lstm_layers = lstm_layers
        self.lstm_hidden = lstm_hidden
        self.bidirectional = bidirectional
        self.dropout = dropout

        net = EfficientNet.from_pretrained('efficientnet-b0',include_top=False )        
        self.cnn = net        
        self.rnn = nn.LSTM(int(1280),
                           self.lstm_hidden, self.lstm_layers,
                           batch_first=True, bidirectional=bidirectional)
        if self.bidirectional:
            self.lin = nn.Linear(2*self.lstm_hidden, 9)
        else:
            self.lin = nn.Linear(self.lstm_hidden, 9)
        if self.dropout:
            self.drop = nn.Dropout(0.5)

    def init_hidden(self, batch_size):
        if self.bidirectional:
            return (Variable(torch.zeros(2*self.lstm_layers, batch_size, self.lstm_hidden).cuda(), requires_grad=True),
                    Variable(torch.zeros(2*self.lstm_layers, batch_size, self.lstm_hidden).cuda(), requires_grad=True))
        else:
            return (Variable(torch.zeros(self.lstm_layers, batch_size, self.lstm_hidden).cuda(), requires_grad=True),
                    Variable(torch.zeros(self.lstm_layers, batch_size, self.lstm_hidden).cuda(), requires_grad=True))

    def forward(self, x, lengths=None):
        batch_size, timesteps, C, H, W = x.size()
        self.hidden = self.init_hidden(batch_size)

        # CNN forward
        c_in = x.view(batch_size * timesteps, C, H, W)
        c_out = self.cnn(c_in)
        c_out = c_out.mean(3).mean(2)
        if self.dropout:
            c_out = self.drop(c_out)

        # LSTM forward
        r_in = c_out.view(batch_size, timesteps, -1)
        r_out, states = self.rnn(r_in, self.hidden)
        out = self.lin(r_out)
        out = out.view(batch_size*timesteps,9)

        return out

위 코드는 CNN-LSTM에서 CNN부분에 EfficientNet을 적용할 때의 코드입니다. 9가지 클래스의 분류문제입니다.

Pytorch Efficientnet코드는 위의 from efficientnet_pytorch import EfficientNet 에서 가져올 수 있습니다.

classifier는 사용하지 않기때문에 include_top=false로 설정해줍니다.

이후 cnn을 efficientNet으로 설정해두고 rnn = LSTM 설정을 해줍니다. 

EfficientNet의 output은 avgpool이기 때문에 1280을 인풋으로 넣어줍니다. 

bidirectional을 적용할 경우에 LSTM layer 개수를 2배 해줍니다. 또한 다음 레이어인 linear에 히든노드를 2배 적용해줍니다.

init에서 모델 파라미터를 설정해주고 forward에서 모델 구조를 설계해줍니다.

x는 인풋 데이터이고, 인풋 데이터는 동영상 데이터입니다. 따라서 batch_first일 경우 위와 같이 분리할 수 있습니다.

각각 CNN과 LSTM에 인풋 데이터를 설정해주고 CNN의 아웃풋을 LSTM에 입력해줍니다.