KOTE_pytorch_lightning 오류 해결

안녕하세요! 이번에는 KOTE_pytorch_lightning 오류 해결에 대한 포스팅을 하겠습니다!

이번 학기에 기계학습 수업을 들으면서 KOTE데이터셋을 이용해야 하는 일이 생겼는데요! 

https://github.com/searle-j/KOTE

GitHub - searle-j/KOTE: Korean Online That-gul Emotions Dataset

 

 

위 깃허브에서 제공한 코드 그대로 실행시키니 아래와 같은 오류가 발생했습니다

NotImplementedError: Support for `training_epoch_end` has been removed in v2.0.0. `KOTETagger` implements this method. You can use the `on_train_epoch_end` hook instead. To access outputs, save them in-memory as instance attributes. You can find migration examples in https://github.com/Lightning-AI/lightning/pull/16520.

 

training_epoch_end가 삭제가 되어서 implement가 안된다는 오류인데요!

여기서 잠깐 모델 구조를 살펴보면

class KOTETagger(pl.LightningModule):
    
    def __init__(self, n_training_steps=None, n_warmup_steps=None, gamma_for_expLR=None):
        super().__init__()
        self.electra = electra
        self.classifier = nn.Linear(self.electra.config.hidden_size, 44) ## the label dimension == 44 <-- what an ominous number for asians though... <-- I didn't intend it!
        self.n_training_steps = n_training_steps
        self.n_warmup_steps = n_warmup_steps
        
        ## the loss
        self.criterion = nn.BCELoss()
        
    def forward(self, input_ids, attention_mask, labels=None):
        output = self.electra(input_ids, attention_mask=attention_mask)
        output = output.last_hidden_state[:,0,:] ## [CLS] of the last hidden state
        output = self.classifier(output)
        output = torch.sigmoid(output)
        loss = 0
        if labels is not None:
            loss = self.criterion(output, labels)
        
        torch.cuda.empty_cache()
        
        return loss, output
    
    def step(self, batch, batch_idx):
        input_ids = batch["input_ids"]
        attention_mask = batch["attention_mask"]
        labels = batch["labels"]
        loss, outputs = self.forward(input_ids, attention_mask, labels)

        preds = outputs

        y_true = list(labels.detach().cpu())
        y_pred = list(preds.detach().cpu())

        return {"loss": loss, "y_true": y_true, "y_pred": y_pred}
    
    def training_step(self, batch, batch_idx):
        return self.step(batch, batch_idx)

    def validation_step(self, batch, batch_idx):
        return self.step(batch, batch_idx)
    
    def epoch_end(self, outputs, state="train"):
        loss = torch.tensor(0, dtype=torch.float)
        for out in outputs:
            loss += out["loss"].detach().cpu()
        loss = loss / len(outputs)

        y_true = []
        y_pred = []
        for out in outputs:
            y_true += out["y_true"]
            y_pred += out["y_pred"]

        self.log(state + "_loss", float(loss), on_epoch=True, prog_bar=True)
        print(f"[Epoch {self.trainer.current_epoch} {state.upper()}] Loss: {loss}")
        return {"loss": loss}
    
    def training_epoch_end(self, outputs):
        self.epoch_end(outputs, state="train")

    def validation_epoch_end(self, outputs):
        self.epoch_end(outputs, state="val")
    
    def configure_optimizers(self):
        optimizer = AdamW(self.parameters(), lr=INITIAL_LR)
        
        scheduler = get_linear_schedule_with_warmup(
            optimizer,
            num_warmup_steps=self.n_warmup_steps,
            num_training_steps=self.n_training_steps
        )
        
        return dict(
          optimizer=optimizer,
          lr_scheduler=dict(
            scheduler=scheduler,
            interval="step"
          )
        )

 

training_epoch_end와 validation_epoch_end class가 있는 것을 확인할 수 있습니다! 오류 해결을 위해 해당 클래스 대신 on_training_epoch_end, on_validation_epoch_end로 바꾸고, 모델 구조에 맞게끔 수정을 진행했습니다!

 

최종 모델 코드는 아래와 같습니다!😀😀

class KOTETagger(pl.LightningModule):
    def __init__(self, n_training_steps=None, n_warmup_steps=None, gamma_for_expLR=None):
        super().__init__()
        self.electra = electra
        self.classifier = nn.Linear(self.electra.config.hidden_size, 44)
        self.n_training_steps = n_training_steps
        self.n_warmup_steps = n_warmup_steps
        self.criterion = nn.BCELoss()

        self.training_step_loss = [] # 훈련 단계의 손실을 저장하는 리스트
        self.val_step_loss = []  #검 단계의 손실을 저장하는 리스트
        self.training_step_outputs = [] #훈련 단계의 출력을 저장하는 리스트
        self.val_step_outputs = [] # 검증 단계의 출력을 저장하는 리스트

    def forward(self, input_ids, attention_mask, labels=None):
        output = self.electra(input_ids, attention_mask=attention_mask) # Electra 모델에 입력을 전달하여 출력 얻기
        output = output.last_hidden_state[:, 0, :]
        output = self.classifier(output)
        output = torch.sigmoid(output) # 시그모이드 함수를 사용하여 출력값을 확률로 변환
        loss = self.criterion(output, labels) if labels is not None else 0
        return loss, output

    def training_step(self, batch, batch_idx):
        input_ids = batch["input_ids"]
        attention_mask = batch["attention_mask"]
        labels = batch["labels"]
        loss, outputs = self.forward(input_ids, attention_mask, labels)

        self.training_step_loss.append(loss.detach().cpu().item())
        self.training_step_outputs.extend([(labels.detach().cpu().numpy(), outputs.detach().cpu().numpy())])
        return {"loss": loss}

    def validation_step(self, batch, batch_idx):
        input_ids = batch["input_ids"]
        attention_mask = batch["attention_mask"]
        labels = batch["labels"]
        loss, outputs = self.forward(input_ids, attention_mask, labels)

        self.val_step_loss.append(loss.detach().cpu().item())
        self.val_step_outputs.extend([(labels.detach().cpu().numpy(), outputs.detach().cpu().numpy())])
        return {"loss": loss}

    def epoch_end(self, state="train"):
      ## 훈련 또는 검증 단계에 대한 손실 평균 계산
        loss = torch.mean(torch.tensor(self.training_step_loss)) if state == "train" else torch.mean(torch.tensor(self.val_step_loss))

        self.log(state + "_loss", loss, on_epoch=True, prog_bar=True)
        print(f"[Epoch {self.trainer.current_epoch} {state.upper()}] Loss: {loss}")
        return {"loss": loss}

    def on_train_epoch_end(self):
        self.epoch_end(state="train")
        self.training_step_loss = []

    def on_validation_epoch_end(self):
        self.epoch_end(state="val")
        self.val_step_loss = []

    def configure_optimizers(self):
        optimizer = AdamW(self.parameters(), lr=INITIAL_LR)
        scheduler = get_linear_schedule_with_warmup(
            optimizer,
            num_warmup_steps=self.n_warmup_steps,
            num_training_steps=self.n_training_steps
        )

        return {
            "optimizer": optimizer,
            "lr_scheduler": {
                "scheduler": scheduler,
                "interval": "step"
            }
        }

 

위 코드에서는 on_train_epoch_end와 on_validation_epoch_end를 사용하여 각 에폭의 끝에서 결과를 처리하도록 변경하였습니다. 

 

> 즉, 기존 코드와 달라진 점을 정리하자면..⚙️⚙️

1. 메서드 이름:
   • 첫 번째 코드: on_train_epoch_end 및 on_validation_epoch_end
   • 두 번째 코드: training_epoch_end 및 validation_epoch_end
2. 출력값의 처리:
   • 첫 번째 코드: on_train_epoch_end와 on_validation_epoch_end에서 outputs를 인자로 받아들이지 않음.
   • 두 번째 코드: training_epoch_end와 validation_epoch_end에서 outputs를 인자로 받아들임.

PyTorch Lightning v1.4.7 이후 버전에서 on_train_epoch_end와 on_validation_epoch_end 메서드가 변경되어 outputs 인자를 사용하지 않게 되었습니다. 대신 training_epoch_end와 validation_epoch_end 메서드에서는 결과를 처리하기 위해 outputs 인자를 사용하도록 설계되었습니다.

두 번째 코드에서는 v1.4.7 이후의 새로운 버전을 따르고 있으며, 새로운 명명 규칙과 인터페이스를 채택하고 있습니다. 이 두 가지 버전의 주요 차이점은 메서드 이름과 메서드가 처리하는 outputs 인자의 유무입니다.

 

> 모델에 대해 간단히 설명드리자면..⚙️⚙️

해당 코드는 PyTorch Lightning을 사용하여 모델을 정의하고 훈련시키는 클래스인 KOTETagger를 나타냅니다.

• __init__: 모델을 초기화하고 필요한 구성 요소 및 하이퍼파라미터를 설정. 또한 Electra 모델과 선형 분류기, 그리고 훈련 및 스케줄링을 위한 하이퍼파라미터들을 설정.
• forward: 주어진 입력 데이터에 대해 모델을 실행하고, Electra 모델을 통과한 후 해당 결과를 선형 분류기에 적용하여 예측값을 생성. 손실 함수를 사용하여 예측된 출력과 실제 레이블 간의 손실을 계산.
• training_step: 훈련 단계에서 배치 데이터를 받아 모델을 실행하고, 훈련 손실을 계산하고 저장.
• validation_step: 검증 단계에서 배치 데이터를 받아 모델을 실행하고, 검증 손실을 계산하고 저장.
• epoch_end: 각 에폭의 끝에서 훈련 및 검증 단계의 손실 평균을 계산하고, 로그에 손실을 기록.
• on_train_epoch_end 및 on_validation_epoch_end: 각각의 에폭이 끝날 때 호출되어 epoch_end 메서드를 호출하고, 손실 리스트를 초기화.
• configure_optimizers: 옵티마이저와 스케줄러를 설정하여 반환.  이때, AdamW 옵티마이저와 선형 스케줄러를 사용하고 있음

 

위와 같은 방식으로 코드를 수정하면 에러가 발생하지 않고, 코드가 잘 실행이 되는 것을 확인할 수 있습니다!

 

 

여기까지가 KOTE데이터셋에 대한 오류 해결과정이었는데요!

혹여나, KOTE 데이터셋을 이용하면서 저와 동일한 오류를 경험한 분들에게 이 글이 조금이나마 도움이 되기를 바라며 해당 포스팅을 마치겠습니다!😀😀