File size: 1,631 Bytes
9067733 03d287b 9067733 |
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 |
import torch
import torch.nn as nn
from DeePixBis.Metrics import test_accuracy, test_loss
class Trainer():
def __init__(self, train_dl, val_dl, model, epochs, opt, loss_fn, device='cpu'):
self.train_dl = train_dl
self.val_dl = val_dl
self.model = model.to(device)
self.epochs = epochs
self.opt = opt
self.loss_fn = loss_fn
self.device = device
def train_one_epoch(self, num):
print(f'\nEpoch ({num+1}/{self.epochs})')
print('----------------------------------')
# self.model.train()
for batch, (img, mask, label) in enumerate(self.train_dl):
img, mask, label = img.to(self.device), mask.to(self.device), label.to(self.device)
net_mask, net_label = self.model(img)
loss = self.loss_fn(net_mask, net_label, mask, label)
# Train
self.opt.zero_grad()
loss.backward()
self.opt.step()
if batch % 9 == 0:
print(f'Loss : {loss}')
# self.model.eval()
test_acc = test_accuracy(self.model, self.val_dl)
test_los = test_loss(self.model, self.val_dl, self.loss_fn)
print(f'Test Accuracy : {test_acc} Test Loss : {test_los}')
return test_acc, test_los
def fit(self):
training_loss = []
training_acc = []
self.model.train()
for epoch in range(self.epochs):
train_acc, train_loss = self.train_one_epoch(epoch)
training_acc.append(train_acc)
training_loss.append(train_loss)
return training_acc, training_loss
|