-
Notifications
You must be signed in to change notification settings - Fork 1
/
batch_engine.py
107 lines (71 loc) · 3.04 KB
/
batch_engine.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
import time
import numpy as np
import torch
from torch.nn.utils import clip_grad_norm_
from tqdm import tqdm
from tools.utils import AverageMeter, to_scalar, time_str
def batch_trainer(epoch, model, train_loader, criterion, optimizer):
model.train()
epoch_time = time.time()
loss_meter = AverageMeter()
batch_num = len(train_loader)
gt_list = []
preds_probs = []
lr = optimizer.param_groups[1]['lr']
for step, (imgs, gt_label, imgname) in enumerate(train_loader):
batch_time = time.time()
imgs, gt_label = imgs.cuda(), gt_label.cuda()
feat_map, output = model(imgs)
loss_list = []
for k in range(len(output)):
out = output[k]
loss_list.append(criterion(out, gt_label))
loss = sum(loss_list)
#maximum voting
output = torch.max(torch.max(torch.max(torch.max(output[0], output[1]), output[2]), output[3]), output[4])
train_loss = loss
optimizer.zero_grad()
train_loss.backward()
clip_grad_norm_(model.parameters(), max_norm=10.0) # make larger learning rate works
optimizer.step()
loss_meter.update(to_scalar(train_loss))
gt_list.append(gt_label.cpu().numpy())
train_probs = torch.sigmoid(output)
preds_probs.append(train_probs.detach().cpu().numpy())
log_interval = 20
if (step + 1) % log_interval == 0 or (step + 1) % len(train_loader) == 0:
print(f'{time_str()}, Step {step}/{batch_num} in Ep {epoch}, {time.time() - batch_time:.2f}s ',
f'train_loss:{loss_meter.val:.4f}')
train_loss = loss_meter.avg
gt_label = np.concatenate(gt_list, axis=0)
preds_probs = np.concatenate(preds_probs, axis=0)
print(f'Epoch {epoch}, LR {lr}, Train_Time {time.time() - epoch_time:.2f}s, Loss: {loss_meter.avg:.4f}')
return train_loss, gt_label, preds_probs
# @torch.no_grad()
def valid_trainer(epoch, model, valid_loader, criterion):
model.eval()
loss_meter = AverageMeter()
preds_probs = []
gt_list = []
with torch.no_grad():
for step, (imgs, gt_label, imgname) in enumerate(tqdm(valid_loader)):
imgs = imgs.cuda()
gt_label = gt_label.cuda()
gt_list.append(gt_label.cpu().numpy())
gt_label[gt_label == -1] = 0
output = model(imgs)
loss_list = []
for k in range(len(output)):
out = output[k]
loss_list.append(criterion(out, gt_label))
loss = sum(loss_list)
# maximum voting
output = torch.max(torch.max(torch.max(torch.max(output[0], output[1]), output[2]), output[3]), output[4])
valid_loss = loss
valid_probs = torch.sigmoid(output)
preds_probs.append(valid_probs.detach().cpu().numpy())
loss_meter.update(to_scalar(valid_loss))
valid_loss = loss_meter.avg
gt_label = np.concatenate(gt_list, axis=0)
preds_probs = np.concatenate(preds_probs, axis=0)
return valid_loss, gt_label, preds_probs