-
Notifications
You must be signed in to change notification settings - Fork 2
/
deep_supervision_solver.py
142 lines (118 loc) · 5.56 KB
/
deep_supervision_solver.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
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
import torch
from torch.optim import Adam
from deep_supervision_model import build_model, weights_init
import numpy as np
import os
import cv2
from loss import TriLoss
# normalize the predicted SOD probability map
def normPRED(d):
ma = torch.max(d)
mi = torch.min(d)
dn = (d - mi) / (ma - mi)
return dn
Dataset_dict = {'e': 'ECSSD', 'p': 'PASCALS', 'd': 'DUTOMRON', 'h': 'HKU-IS', 's': 'SOD', 't': 'DUTS_TE'}
class Solver(object):
def __init__(self, train_loader, test_loader, config):
self.train_loader = train_loader
self.test_loader = test_loader
self.config = config
self.iter_size = config.iter_size
self.show_every = config.show_every
self.lr_decay_epoch = [20, ]
self.loss = TriLoss(sigma=self.config.sigma, n=self.config.weight)
self.build_model()
if config.mode == 'test':
print('Loading pre-trained model from %s...' % self.config.model)
if self.config.cuda:
self.net.load_state_dict(torch.load(self.config.model))
else:
self.net.load_state_dict(torch.load(self.config.model, map_location='cpu'))
self.net.eval()
# print the network information and parameter numbers
def print_network(self, model, name):
num_params = 0
for p in model.parameters():
num_params += p.numel()
print(name)
print(model)
print("The number of parameters: {}".format(num_params))
def deep_supervision_loss_old(self, preds, gt):
losses = []
for pred in preds:
wbce_, iou_, rev_iou_ = self.loss(pred, gt)
losses.append((wbce_ + 0.5 * (iou_ + rev_iou_)).mean())
sum_loss = losses[0] + losses[1] / 2 + losses[2] / 4 + losses[3] / 8 + losses[4] / 8
return sum_loss
def deep_supervision_loss(self, preds, gt):
losses = []
for pred in preds:
wbce_, iou_, rev_iou_ = self.loss(pred, gt)
losses.append((wbce_ + 0.5 * (iou_ + rev_iou_)).mean())
sum_loss = losses[0] + losses[1] / 2 + losses[2] / 4 + losses[3] / 8 + losses[4] / 8
return sum_loss
# build the network
def build_model(self):
self.net = build_model(self.config.arch)
if self.config.cuda:
self.net = self.net.cuda()
self.net.train()
self.net.apply(weights_init)
if self.config.load == '':
self.net.base.load_state_dict(torch.load(self.config.pretrained_model))
else:
self.net.load_state_dict(torch.load(self.config.load))
self.lr = self.config.lr
self.wd = self.config.wd
param1 = [p for p in self.net.parameters() if p.requires_grad]
param2 = [p for p in self.loss.parameters() if p.requires_grad]
param = param1 + param2
self.optimizer = Adam(param, lr=self.lr, weight_decay=self.wd)
def test(self):
self.net.eval()
for i, data_batch in enumerate(self.test_loader):
images, name, im_size = data_batch['image'], data_batch['name'][0], np.asarray(data_batch['size'])
im_size = im_size[1], im_size[0]
# print(im_size)
with torch.no_grad():
if self.config.cuda:
images = images.cuda()
preds = self.net(images)[0]
preds = normPRED(torch.sigmoid(preds))
pred = np.squeeze(preds.cpu().data.numpy())
multi_fuse = 255 * pred
multi_fuse = cv2.resize(multi_fuse, dsize=im_size, interpolation=cv2.INTER_LINEAR)
cv2.imwrite(os.path.join(self.config.test_fold, name[:-4] + '.png'), multi_fuse)
print(Dataset_dict[self.config.sal_mode] + ' Test Done!')
# training phase
def train(self):
iter_num = len(self.train_loader.dataset) // self.config.batch_size
for epoch in range(self.config.epoch):
self.net.zero_grad()
for i, data_batch in enumerate(self.train_loader):
sal_image, sal_label = data_batch['sal_image'], data_batch['sal_label']
if (sal_image.size(2) != sal_label.size(2)) or (sal_image.size(3) != sal_label.size(3)):
print('IMAGE ERROR, PASSING```')
continue
if self.config.cuda:
# cudnn.benchmark = True
sal_image, sal_label = sal_image.cuda(), sal_label.cuda()
sal_preds = self.net(sal_image)
sum_loss = self.deep_supervision_loss(sal_preds, sal_label)
self.optimizer.zero_grad()
sum_loss.backward()
self.optimizer.step()
if i % (self.show_every // self.config.batch_size) == 0:
print(
'epoch: [%2d/%2d], iter: [%5d/%5d] || sum_loss : %10.4f' % (
epoch, self.config.epoch, i, iter_num, sum_loss.data))
print('Learning rate: ' + str(self.lr))
if (epoch + 1) % self.config.epoch_save == 0:
torch.save(self.net.state_dict(), '%s/models/epoch_%d.pth' % (self.config.save_folder, epoch + 1))
if epoch in self.lr_decay_epoch:
self.lr = self.lr * 0.1
param1 = [p for p in self.net.parameters() if p.requires_grad]
param2 = [p for p in self.loss.parameters() if p.requires_grad]
param = param1 + param2
self.optimizer = Adam(param, lr=self.lr, weight_decay=self.wd)
torch.save(self.net.state_dict(), '%s/models/final.pth' % self.config.save_folder)