-
Notifications
You must be signed in to change notification settings - Fork 31
/
train_mc_npair.py
187 lines (162 loc) · 7.13 KB
/
train_mc_npair.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
import torch
import torch.optim as optim
from data import EbayDataset
import os
from configs import opt
from models import ModGoogLeNet,NpairLoss
from torch.utils.data import DataLoader
from tqdm import tqdm
from utils import Visulizer
import csv
import numpy as np
def train(**kwargs):
print("run train")
opt.parse(kwargs)
os.environ['CUDA_VISIBLE_DEVICES'] = str(opt.gpu_id)
if opt.debug:
import ipdb
ipdb.set_trace()
model =ModGoogLeNet(embeding_size=opt.embeding_size)
if opt.dml_model_path:
model.load(opt.dml_model_path)
if opt.use_gpu:
model = model.cuda()
#model.freeze_model(level=opt.freeze_level)
if opt.use_viz:
viz = Visulizer(host=opt.vis_host,port=opt.vis_port,env='dml'+opt.vis_env)
viz.log("start to train dml npair mc model")
#loss function
criterion = NpairLoss(l2_reg=opt.l2_reg)
lr = opt.lr
m = opt.momentum
optimizer = optim.SGD([{'params':model.level1_2.parameters()},
{'params': model.level_3_4.parameters()},
{'params': model.level_5_6.parameters()},
{'params': model.level_7.parameters()},
{'params':model.fc.parameters(),'lr':10*lr}],lr=lr,momentum=m)
#optimizer = optim.SGD(model.parameters(),lr=lr,momentum=m)
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,1,0.2)
# data and dataloader
train_data = EbayDataset(dir_root=opt.ebay_dir, train=True, batch_size=opt.batch_size)
cycle_length = len(train_data)
#val_data = EbayDataset(dir_root=opt.ebay_dir, train=False, batch_size=opt.batch_size)
train_dataloader = DataLoader(train_data, batch_size=1, shuffle=True, num_workers=opt.num_workers)
#val_dataloader = DataLoader(val_data, batch_size=60, shuffle=False, num_workers=opt.num_workers)
print("dataloader setted ,begin to train")
#f = open('dml_log.out','w')
for epoch in range(opt.max_epoch):
lr_scheduler.step()
train_loss = 0
for i,data in enumerate(train_dataloader):
# if i in [200, 800, 1500]:
# lr_scheduler.step()
data = data[0]
if opt.use_gpu:
data = data.cuda()
optimizer.zero_grad()
feature = model(data)
batch_size = data.size(0)
target = torch.arange(0, int(batch_size / 2), dtype=torch.int64).cuda()
loss = criterion(feature,target)
loss.backward()
optimizer.step()
train_loss += loss.item()
freq = int(opt.print_freq)
if i%freq==(freq-1):
average_loss = train_loss /opt.print_freq
#f.write("iteration:{0},dml_loss:{1}\n".format(i+ epoch*cycle_length,average_loss))
#f.flush()
if opt.use_viz:
viz.plot('dml_loss',average_loss)
train_loss =0
if opt.debug:
break
#f.write("epoch:{0} finished,begin to valid test".format(epoch))
model.save()
# if epoch>1 and epoch%5==0:
# val(model,val_dataloader,epoch)
if opt.debug:
#f.write("finish one iter")
break
#f.write("finish train epoch {0}".format(opt.max_epoch))
#f.close()
def val(model,dataloder,epoch):
"""
this val model will calculate the nmi index.normal mutual information
:param model: the emebding model
:param dataloder: val dataloder
:return:
"""
# prepare file model to extract feature
file_name = 'checkpoints/online_product_{0}.csv'.format(epoch)
f = open(file_name,'w')
writer = csv.writer(f,dialect='excel')
model.eval()
# feature extreat,fisrt for all image,image_id,class_id extract the feature vector
for i,(data,image_id,class_id) in enumerate(dataloder):
if opt.use_gpu:
data = data.cuda()
feature = model(data) # the feature is [batch,512] vector
vector = feature.cpu().detach().numpy() if opt.use_gpu else feature.numpy()
image_id = image_id.numpy().reshape(-1,1)
class_id = class_id.numpy().reshape(-1,1)
result = np.hstack(image_id,class_id,vector)
#write the data to dataframe file
writer.writerows(result)
if opt.debug:
print("test one batch of val data and save to csv file")
break
f.close()
# clustering to centriod, second, for all image_id,and corresponding feature vector,using kmeans cluster to fixed |class_id|
#featuredata = pd.read_csv(file_name,header=None)
# for distribution from origin and cluster distribution.compute the nmi by sklearn metric nmi
model.train()
print("finished cluster and evalution")
def compute(**kwargs):
print("run compute_vector")
opt.parse(kwargs)
os.environ['CUDA_VISIBLE_DEVICES'] = str(opt.gpu_id)
if opt.debug:
import ipdb
ipdb.set_trace()
model =ModGoogLeNet(embeding_size=opt.embeding_size)
if opt.dml_model_path:
model.load(opt.dml_model_path)
if opt.use_gpu:
model = model.cuda()
val_data = EbayDataset(dir_root=opt.ebay_dir, train=False, batch_size=opt.batch_size)
val_dataloader = DataLoader(val_data, batch_size=opt.batch_size, shuffle=False, num_workers=opt.num_workers)
file_name = 'checkpoints/online_product_compute.csv'
f = open(file_name, 'w')
writer = csv.writer(f, dialect='excel')
model.eval()
# feature extreat,fisrt for all image,image_id,class_id extract the feature vector
for i, (data, image_id, class_id) in enumerate(val_dataloader):
if opt.use_gpu:
data = data.cuda()
feature = model(data) # the feature is [batch,512] vector
vector = feature.cpu().detach().numpy() if opt.use_gpu else feature.numpy()
image_id = image_id.numpy().reshape(-1, 1)
class_id = class_id.numpy().reshape(-1, 1)
result = np.hstack([image_id, class_id, vector])
# write the data to dataframe file
writer.writerows(result)
if opt.debug:
print("test one batch of val data and save to csv file")
break
f.close()
def help():
"""print function use information"""
print("""this file help to train product train:
exanple --:
python train_mc_npair.py help
python train_mc_npair.py train --gpu_id=3 --debug=True
python train_mc_npair.py train --gpu_id=2 --batch_size=72
python train_mc_npair.py train --gpu_id=3 --lr=0.0003 --batch_size=72
python train_mc_npair.py train --gpu_id=0 --debug=True --dml_model_path=checkpoints/DMLGoogle_0710_20\:24\:04.pth
python train_mc_npair.py train --batch_size=120 --gpu_id=3 --lr=0.0001 --debug=True --dml_model_path=checkpoints/DMLGoogle_0710_20\:24\:04.pth
python train_mc_npair.py compute --batch_size=300 --gpu_id=2 --dml_model_path=checkpoints/DMLGoogle_0714_07:51:44.pth --num_workers=6
""")
if __name__=='__main__':
import fire
fire.Fire()