모델 학습 과정에서, 기존의 손실 함수가 아닌 다른 손실 함수를 사용하고 싶을 경우가 있다.
이럴 때는 아래 예시와 같은 방법을 사용해보자.
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()