-
Notifications
You must be signed in to change notification settings - Fork 0
/
trainer.py
79 lines (52 loc) · 2.35 KB
/
trainer.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
import torch
import torch.nn as nn
import torch.nn.functional as F
import pandas as pd
from torchtext.legacy.data.iterator import batch
#%%
class Trainer(object):
def __init__(self, args, loaders, model):
self.args = args
self.model = model
self.loaders = loaders
trg_pad_idx = self.loaders.trg.vocab.stoi[self.loaders.trg.pad_token]
self.device = torch.device('cuda')
self.model.to(self.device)
self.criterion = nn.CrossEntropyLoss(ignore_index = trg_pad_idx).to(self.device)
self.optimizer = torch.optim.Adam(self.model.parameters())
def train(self, data_loader):
epoch_loss= 0
self.model.train() # train mode
for i, batch in enumerate(data_loader): # index, data # 위에 출력한거 확인! i=train loader의 길이-1
src=batch.src
trg=batch.trg
self.optimizer.zero_grad()
output=self.model(src,trg)
#trg = [trg len, batch size]
#output = [trg len, batch size, output dim]
output_dim=output.shape[-1]
# loss 함수는 2d input으로만 계산 가능
output = output[1:].view(-1, output_dim)
trg = trg[1:].view(-1)
# trg = [(trg len-1) * batch size]
# output = [(trg len-1) * batch size, output dim)]
loss = self.criterion(output, trg)
loss.backward()
self.optimizer.step()
epoch_loss+=loss.item()
return epoch_loss/len(data_loader)
def evaluate(self, data_loader):
epoch_loss=0
self.model.eval() # evaluation mode
with torch.no_grad(): #backprop 기능 끄기
for i, batch in enumerate(data_loader):
src = batch.src
trg = batch.trg
# teacher_forcing_ratio = 0 (아무것도 알려주면 안 됨)
output = self.model(src, trg, 0)
output_dim = output.shape[-1]
output = output[1:].view(-1, output_dim)
trg = trg[1:].view(-1)
loss = self.criterion(output, trg)
epoch_loss+=loss.item()
return epoch_loss/len(data_loader)