[PyTorch] nn.Module로 모델 custom하기!!

지난 번에는 BertModel과 BertForSequenceClassification에 대한 비교를 다루었습니다.

https://chanmuzi.tistory.com/243

[PyTorch] AutoModel vs AutoModelForSequenceClassification 비교하기 (BERT 파헤치기!!)

 

입력은 동일하지만 출력이 다르다는 사실, 그리고 그 이유를 코드와 함께 자세히 설명했었는데요,

이번에는 nn.Module을 이용하여 BertModel을 BertForSequenceClassifcaion으로 만들어보겠습니다.

 

결론적으로 이야기하자면 nn.Module은 크게 보면,

1. 바닥(scratch)부터 딥러닝 모델을 짜거나

2. 기존의 모델에 추가적인 변형을 주는 등의 작업

을 할 수 있습니다.

---

classifier를 달아보자  🔥🔥

지난 포스팅에서 다루었던 내용으로, 왼쪽이 BertModel, 오른쪽이 BertForSequenceClassification입니다.

다른 것은 다 동일하지만 BertPooler 이후 Dropout, Linear가 추가로 존재하는지 여부에서 차이가 있음을 알 수 있습니다.

오늘은 저 부분을 직접 구현해보려고 해요.

 

from transformers import BertModel, BertForSequenceClassification

basic = BertModel.from_pretrained('bert-base-uncased')
sequence = BertForSequenceClassification.from_pretrained('bert-base-uncased')

우선 두 모델을 불러옵니다.

이 상태에서 각각의 변수를 출력하면 위의 이미지같은 결과가 나타납니다.

(이건 그냥 불러오는 걸 보여드리려고 작성한 코드에요)

 

import torch.nn

nn.Module을 불러오기 위해 torch 라이브러리를 불러옵니다.

 

class basic_sequence(nn.Module):
    def __init__(self):
        super(basic_sequence, self).__init__()
        self.bert = BertModel.from_pretrained('bert-base-uncased')
        self.dropout = nn.Dropout(p=0.1, inplace=False)
        self.classifier = nn.Linear(in_features=768, out_features=2, bias=True)

    def forward(self, input_ids, attention_mask, labels=None):
        outputs = self.bert(input_ids,attention_mask)
        pooled_output = outputs[1]
        pooled_output = self.dropout(pooled_output)
        logits = self.classifier(pooled_output)

        if labels is not None:
            loss_fct = nn.CrossEntropyLoss()
            loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
            return loss
        else:
            return logits

갑자기 이해하기 힘든 내용들이 훅 들어온 것 같네요..!

하나씩 살펴봅시다.

---

nn.Module은 크게 두 부분으로 나누어 사용합니다.

init과 forward로 말이죠.

init은 상속받는 모듈의 기능을 가져오겠다는 의미이고, forward는 변수들을 말 그대로 보내겠다는 뜻으로 이해할 수 있습니다.

    def __init__(self):
        super(basic_sequence, self).__init__()
        self.bert = BertModel.from_pretrained('bert-base-uncased')
        self.dropout = nn.Dropout(p=0.1, inplace=False)
        self.classifier = nn.Linear(in_features=768, out_features=2, bias=True)

 

__init__(self) 🔥

첫 줄의 super가 바로 상속을 의미하는 코드죠.

부모 클래스의 메서드를 불러오는 역할입니다.

 

다음으로 bert 모델을 불러옵니다.

이 방식은 우리가 class 바깥에서 모델의 가중치를 불러오던 것과 동일한 방식이에요.

하지만 class의 인스턴스가 생기는 순간에 이 모델을 불러오도록 설정해준 것입니다.

 

다음으로 dropout, linear가 있는데요, 이 둘은 이후 forward 함수에서 사용될 함수들입니다.

즉, init에는 '상속'과 'forward를 위한 세팅'을 담당하는 코드들이 포함됩니다.

만약 내가 이 모델을 선언하면서 반드시 포함시키고 싶은 변수들이 있다(이 변수가 빠지면 에러가 발생하게 만들고싶다), 하는 경우에도 이 init에 변수를 포함시키면 됩니다.

---

forward 🔥🔥

모델에 어떤 입력을 주면 이것이 특정 계산에 의해 출력되고, 이 과정을 forward라 할 수 있겠죠.

forward 함수 내에는 어떤 식으로 계산을 할지를 정해주는 코드들이 작성됩니다.

 

    def forward(self, input_ids, attention_mask, labels=None):
        outputs = self.bert(input_ids,attention_mask)
        pooled_output = outputs[1]
        pooled_output = self.dropout(pooled_output)
        logits = self.classifier(pooled_output)

        if labels is not None:
            loss_fct = nn.CrossEntropyLoss()
            loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
            return loss
        else:
            return logits

우선 init에 선언한 bert 모델에 입력을 주어 output을 구해줍니다.

(bert 모델의 output은 'last_hidden_state'와 'pooler_output'으로 구성되어 있습니다)

그리고 첫 번째 값을 꺼내어 pooler_output을 가져오죠.

여기에 dropout을 적용하고 classifier(분류기)에 입력으로 제공합니다.

 

init에서 classifier는 nn.Linear 함수로 구현되었기 때문에, 사실상 pooler_output에 dropout을 적용하고 nn.Linear를 적용했다, 고 정리할 수 있습니다.

 

나머지 label에 관한 코드는 우리가 모델에 입력으로 주는 것이 training/validation dataset인지 test dataset인지에 따라 label 유무가 달라지기 때문에 구분하여 코드를 작성한 것입니다.

만약 label이 존재한다면 바로 loss를 계산하여 반환하고, 그렇지 않은 경우에는 예측 결과인 logits만 반환하면 되겠죠.

 

 

이제 class를 호출하여 instance를 생성하고 이를 출력해보면..

basic = basic_sequence()
basic

이렇게 BertForSequenceClassification과 동일한 architecture를 가지게 되었음을 확인할 수 있습니다.

---

맨 처음 설명에서, nn.Module을 이용하여 밑바닥부터 딥러닝 모델을 짤 수 있다고도 말씀드렸는데요,

위에 이미지에서 보이는 각 줄을 이루는 함수나 클래스들을 불러와서 init, forward에 포함시키면 됩니다.

 

또는 지금처럼 사전학습된 모델을 불러온 다음 classifier을 달아주거나, LSTM 등과 같은 legacy 모델을 추가하거나, attention mechanism을 적용해보는 등의 추가작업을 해볼 수 있습니다.

 

각자가 다루고 있는 태스크에 따라 입맞에 맞게 모델을 다뤄서 성능 향상을 꾀하는 시도를 해보는 것이 중요한 것 같습니다.

그리고 뭐니뭐니해도 본인이 직접 document를 보면서 어떻게 구성되어 있는지 파악하는 것이 제일 우선인 것 같구요.

(https://pytorch.org/docs/stable/generated/torch.nn.Module.html)

(https://github.com/huggingface/transformers/blob/v4.27.2/src/transformers/models/bert/modeling_bert.py#L873)

 

 

혹시라도 잘못되었거나 부족한 내용에 대해 지적해주시면 대단히 감사하겠습니다 🙇‍♂️