-
Notifications
You must be signed in to change notification settings - Fork 9
/
Copy pathtrain_utils.py
125 lines (95 loc) · 3.97 KB
/
train_utils.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
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
import torch
import torch.nn as nn
from torch.autograd import Variable
import torch.nn.functional as F
from data import NoisyDataset
import torch.optim as optim
from torch.optim import lr_scheduler
from tqdm import tqdm
from torch.utils.data import DataLoader
class Train:
def __init__(self, architecture, train_dir, val_dir, params):
self.cuda = params['cuda']
if self.cuda:
self.architecture = architecture.cuda()
else:
self.architecture = architecture
self.train_dir = train_dir
self.val_dir = val_dir
self.noise_model = params['noise_model']
self.crop_size = params['crop_size']
self.clean_targs = params['clean_targs']
self.lr = params['lr']
self.epochs = params['epochs']
self.bs = params['bs']
self.train_dl, self.val_dl = self.__getdataset__()
self.optimizer = self.__getoptimizer__()
self.scheduler = self.__getscheduler__()
self.loss_fn = self.__getlossfn__(params['lossfn'])
def train(self):
for _ in range(self.epochs):
tr_loss = 0
self.architecture.train()
for _list in tqdm(self.train_dl):
if self.cuda:
source = _list[0].cuda()
target = _list[-1].cuda()
else:
source = _list[0]
target = _list[-1]
_op = self.architecture(Variable(source))
if len(_list) == 4:
if self.cuda:
mask = Variable(_list[1].cuda())
else:
mask = Variable(_list[1])
_loss = self.loss_fn(mask * _op, mask * Variable(target))
else:
_loss = self.loss_fn(_op, Variable(target))
tr_loss += _loss.data
self.optimizer.zero_grad()
_loss.backward()
self.optimizer.step()
val_loss = self.evaluate()
#self.scheduler.step(val_loss)
print(f'Training loss = {tr_loss}, Validation loss = {val_loss}')
def evaluate(self):
val_loss = 0
self.architecture.eval()
for _, _list in enumerate(self.val_dl):
if self.cuda:
source = _list[0].cuda()
target = _list[-1].cuda()
else:
source = _list[0]
target = _list[-1]
_op = self.architecture(Variable(source))
if len(_list) == 4:
if self.cuda:
mask = Variable(_list[1].cuda())
else:
mask = Variable(_list[1])
_loss = self.loss_fn(mask * _op, mask * Variable(target))
else:
_loss = self.loss_fn(_op, Variable(target))
val_loss += _loss.data
return val_loss
def __getdataset__(self):
train_ds = NoisyDataset(self.train_dir, crop_size=self.crop_size, train_noise_model=self.noise_model,
clean_targ=self.clean_targs)
train_dl = DataLoader(train_ds, batch_size=self.bs, shuffle=True)
val_ds = NoisyDataset(self.val_dir, crop_size=self.crop_size, train_noise_model=self.noise_model,
clean_targ=True)
val_dl = DataLoader(val_ds, batch_size=self.bs)
return train_dl, val_dl
def __getoptimizer__(self):
return optim.Adam(self.architecture.parameters(), self.lr)
def __getscheduler__(self):
return lr_scheduler.ReduceLROnPlateau(self.optimizer, patience=self.epochs/4, factor=0.5, verbose=True)
def __getlossfn__(self, lossfn):
if lossfn == 'l2':
return nn.MSELoss()
elif lossfn == 'l1':
return nn.L1Loss()
else:
raise ValueError('No such loss function supported')