[PyTorch] Custom Loss Function 사용하기

모델 학습 과정에서, 기존의 손실 함수가 아닌 다른 손실 함수를 사용하고 싶을 경우가 있다.

이럴 때는 아래 예시와 같은 방법을 사용해보자.

Custom Loss Function

import torch

def MyLossFunction(Y, Y_pred):
	loss = Y - Y_pred
    loss_tensor = torch.tensor(loss, requires_grad=True)
    
    # GPU 사용 시
    loss_tensor = torch.tensor(loss, requires_grad=True).to('device')
    
    return loss_tensor

학습 시에 GPU를 사용하기 위해서는 모델뿐만 아니라 텐서도 GPU에 올려야한다는 것을 기억하자.

학습 진행하기

from transformers import AdamW

NUM_EPOCHS = 10

optimizer = AdamW(model.parameters(), lr=1e-5)
model.train()		# train mode로 바꿔주기

for epoch in range(NUM_EPOCHS):
	for batch in MyDataLoader:
    	x = batch['input']
        Y = batch['output']
        Y_pred = model(Y)		# model에 따라 실행 방법이 다르니 유의하기
        loss_out = MyLossFunction(Y, Y_pred)
        
        optimizer.zero_grad()
        loss_out.backward()
        optimizer.step()