-
Notifications
You must be signed in to change notification settings - Fork 31
/
train_margin_cub.py
199 lines (171 loc) · 7.74 KB
/
train_margin_cub.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
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
import argparse
import time
import logging
import numpy as np
import torch
from torch.utils.data import DataLoader
import os
from data import CUB200DataSet
from models import Margin_Loss,SampleModel
parser = argparse.ArgumentParser(description="train a margin based loss model")
parser.add_argument('--data_path',type=str,default="data/cub200_2011",
help='path of the cub_data')
parser.add_argument('--embed_dim',type=int,default=128,
help='dimensionality of image embeding,times of 8')
parser.add_argument('--batch_size',type=int,default=70,
help='training batch size per device')
parser.add_argument('--batch_k',type=int,default=5,
help='number of images per class in a batch,can be divided by batch_size')
parser.add_argument('--gpu_id',type=str,default='0',
help='the gpu_id of the runing batch')
parser.add_argument('--epochs',type=int,default=100,
help='number of training epochs,default is 100')
parser.add_argument('--optimizer',type=str,default='adam',
help='optimizer,default is adam')
parser.add_argument('--lr',type=float,default=0.0001,
help='learning rate of the resnet and dense layer')
parser.add_argument('--lr_beta',type=float,default=0.1,
help='learning rate for the beta in margin based loss')
parser.add_argument('--margin',type=float,default=0.2,
help='margin for the margin based loss,default is 0.2')
parser.add_argument('--beta',type=float,default=1.2,
help='the class specific beta parameter')
parser.add_argument('--nu',type=float,default=0.0,
help='regularization parameter for beta,default is 0')
parser.add_argument('--steps',type=str,default='30,50,100,300',
help='epochs to updata learning rate')
parser.add_argument('--wd',type=float,default=0.0001,
help='weight decay rate,default is 0.0001')
parser.add_argument('--seed',type=int,default=123,
help='random seed to use,default=123')
parser.add_argument('--factor',type=float,default=0.5,
help='learning rate schedule factor,default is 0.5')
parser.add_argument('--print_freq',type=int,default=20,
help='print the accumulate loss for training process')
parser.add_argument('--debug',action='store_true',default=False)
opt = parser.parse_args()
logging.info(opt)
torch.random.manual_seed(opt.seed)
np.random.seed(opt.seed)
batch_size = opt.batch_size
os.environ['CUDA_VISIBLE_DEVICES']=opt.gpu_id
steps = [int(step) for step in opt.steps.split(',')]
def train():
"""
train the margin based loss model
:return:
"""
# prepare for data for loader
train_data = CUB200DataSet(data_path='data/cub200_2011/CUB_200_2011',batch_k=opt.batch_k,batch_size = opt.batch_size,is_train=True)
test_data = CUB200DataSet(data_path='data/cub200_2011/CUB_200_2011',is_train=False)
train_loader = DataLoader(train_data,batch_size=1,shuffle=False,num_workers=6)
test_loader = DataLoader(test_data,batch_size=60,shuffle=False,num_workers=6)
#begin to set model loss,optimizer,lr_rate, lr_schedule
model = SampleModel(embeding_dim=opt.embed_dim)
beta = torch.tensor(np.ones(100)*opt.beta, requires_grad=True,dtype=torch.float32)
loss_criterion = Margin_Loss(batch_k=opt.batch_k,\
margin=opt.margin,nu=opt.nu) # set loss function for this model
conv_params = []
non_conv_param =[]
for name,param in model.base_model.named_parameters():
if 'conv' in name:
conv_params.append({'params':param,'lr':opt.lr*0.01})
else:
non_conv_param.append({'params':param,'lr':opt.lr})
total_param =[]
total_param.append({'params':model.dense.parameters(),'lr':opt.lr})
total_param.extend(conv_params)
total_param.extend(non_conv_param)
#optimizer = torch.optim.Adam(total_param,lr=opt.lr,weight_decay=opt.wd)
optimizer = torch.optim.SGD(total_param,lr=opt.lr,momentum=0.89)
optimizer_beta = torch.optim.SGD([{'params':beta}],lr=opt.lr_beta,momentum= 0.9)
lr_schedule = torch.optim.lr_scheduler.MultiStepLR(optimizer,
milestones=steps,gamma=opt.factor)
if int(opt.gpu_id)>=0:
model = model.cuda() # the loss function has paramter to convey to cuda
beta = beta.cuda() # the beta parameter has parameter to stored in cuda
loss_criterion = loss_criterion.cuda() # the loss criterion has compute in cuda
loss_criterion.convert_param(to_cuda=True)
# begin to fetch data and train model
for epoch in range(opt.epochs):
print("begin to train epochs:{0}",epoch)
cumulative_loss =0
prev_loss = 0
lr_schedule.step()
for i,data in enumerate(train_loader):
images,label = data[0][0],data[1][0]
if int(opt.gpu_id)>=0:
images = images.cuda()
label = label.cuda()
features = model(images)
loss = loss_criterion(features,label,beta)
loss.backward()
optimizer.step()
optimizer_beta.step()
cumulative_loss += loss.item()
if (i+1)%(opt.print_freq)==0:
print("[Epoch %d,Iter %d] training loss=%f"%(epoch,i+1,cumulative_loss-prev_loss))
prev_loss = cumulative_loss
if opt.debug:
break
print("[Epoch %d] trainin loss =%f"%(epoch,cumulative_loss))
# print test val recall index
names,val_accs = val_model(model,test_loader)
for name,val_acc in zip(names,val_accs):
print("Epoch %d,validation:%s=%f"%(epoch,name,val_acc))
print("job finished")
def val_model(model,test_loader):
"""
val the model,return the recall@K k=1 index
:param model: Margin based model to extract feature of 128 dimension
:param test_loader: Test dataloader to load images data
:return: the recall@K k=1 index
"""
model.eval()
outputs = []
labels =[]
with torch.no_grad():
for data,label in test_loader:
if int(opt.gpu_id)>=0:
data = data.cuda()
feature = model(data)
outputs += feature.detach().cpu().numpy().tolist()
labels += label.numpy().tolist()
model.train()
#eval recall@k
features = np.array(outputs)
labels = np.array(labels)
return evaluate_emb(features,labels)
def evaluate_emb(features,labels):
"""
evaluate embedding in recall
:param features:
:param labels:
:return:
"""
d_mat = get_distance_matrix(features)
names =[]
accs =[]
for k in [1,2,4,8,16]:
names.append('Recall@%d'%k)
correct,cnt = 0.0,0.0
for i in range(features.shape[0]):
d_mat[i,i]=1e10
nns = d_mat[i].argpartition(k)[:k]
if any(labels[i] ==labels[nn] for nn in nns):
correct +=1
cnt +=1
accs.append(correct/cnt)
return names,accs # names is a list of ["Recall@K",,,,] accs is a list of [float_value]
def get_distance_matrix(x):
"""
compute the distance matirx of features,
:param x: np.ndarray in shape (n,d) d is 128
:return: distance matrix of [n,n] for distance in each vector
"""
squrare = np.sum(x*x,axis=1,keepdims=True)
distance_squrare = squrare + squrare.transpose() -2*np.dot(x,x.transpose())
return distance_squrare
if __name__=='__main__':
print("begin to train the model of margin based loss")
train()