-
Notifications
You must be signed in to change notification settings - Fork 33
/
utils.py
132 lines (98 loc) · 3.83 KB
/
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
125
126
127
128
129
130
131
132
import torchvision
import os
import torch
from pylab import plt
def load_resnet50_param(model, stop_layer='layer4'):
resnet50 = torchvision.models.resnet50(pretrained=True)
saved_state_dict = resnet50.state_dict()
new_params = model.state_dict().copy()
for i in saved_state_dict: # copy params from resnet50,except layers after stop_layer
i_parts = i.split('.')
if not i_parts[0] == stop_layer:
new_params['.'.join(i_parts)] = saved_state_dict[i]
else:
break
model.load_state_dict(new_params)
model.train()
return model
def check_dir(checkpoint_dir):#create a dir if dir not exists
if not os.path.exists(checkpoint_dir):
os.makedirs(os.path.join(checkpoint_dir,'model'))
os.makedirs(os.path.join(checkpoint_dir,'pred_img'))
def optim_or_not(model, yes):
for param in model.parameters():
if yes:
param.requires_grad = True
else:
param.requires_grad = False
def turn_off(model):
optim_or_not(model.module.conv1, False)
optim_or_not(model.module.layer1, False)
optim_or_not(model.module.layer2, False)
optim_or_not(model.module.layer3, False)
def get_10x_lr_params(model):
"""
get layers for optimization
"""
b = []
b.append(model.module.layer5.parameters())
b.append(model.module.layer55.parameters())
b.append(model.module.layer6_0.parameters())
b.append(model.module.layer6_1.parameters())
b.append(model.module.layer6_2.parameters())
b.append(model.module.layer6_3.parameters())
b.append(model.module.layer6_4.parameters())
b.append(model.module.layer7.parameters())
b.append(model.module.layer9.parameters())
b.append(model.module.residule1.parameters())
b.append(model.module.residule2.parameters())
b.append(model.module.residule3.parameters())
for j in range(len(b)):
for i in b[j]:
yield i
def loss_calc_v1(pred, label, gpu):
label = label.long()
criterion = torch.nn.CrossEntropyLoss(ignore_index=255).cuda(gpu)
return criterion(pred, label)
def plot_loss(checkpoint_dir,loss_list,save_pred_every):
x=range(0,len(loss_list)*save_pred_every,save_pred_every)
y=loss_list
plt.switch_backend('agg')
plt.plot(x,y,color='blue',marker='o',label='Train loss')
plt.xticks(range(0,len(loss_list)*save_pred_every+3,(len(loss_list)*save_pred_every+10)//10))
plt.legend()
plt.grid()
plt.savefig(os.path.join(checkpoint_dir,'loss_fig.pdf'))
plt.close()
def plot_iou(checkpoint_dir,iou_list):
x=range(0,len(iou_list))
y=iou_list
plt.switch_backend('agg')
plt.plot(x,y,color='red',marker='o',label='IOU')
plt.xticks(range(0,len(iou_list)+3,(len(iou_list)+10)//10))
plt.legend()
plt.grid()
plt.savefig(os.path.join(checkpoint_dir,'iou_fig.pdf'))
plt.close()
def get_iou_v1(query_mask,pred_label,mode='foreground'):#pytorch 1.0 version
if mode=='background':
query_mask=1-query_mask
pred_label=1-pred_label
num_img=query_mask.shape[0]#batch size
num_predict_list,inter_list,union_list,iou_list=[],[],[],[]
for i in range(num_img):
num_predict=torch.sum((pred_label[i]>0).float()).item()
combination = (query_mask[i] + pred_label[i]).float()
inter = torch.sum((combination == 2).float()).item()
union = torch.sum((combination ==1).float()).item()+torch.sum((combination ==2).float()).item()
if union!=0:
inter_list.append(inter)
union_list.append(union)
iou_list.append(inter/union)
num_predict_list.append(num_predict)
else:
inter_list.append(inter)
union_list.append(union)
iou_list.append(0)
num_predict_list.append(num_predict)
return inter_list,union_list,iou_list,num_predict_list