-
Notifications
You must be signed in to change notification settings - Fork 2
/
train.py
executable file
·197 lines (154 loc) · 8.09 KB
/
train.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
# coding=utf-8
from __future__ import absolute_import, print_function
import time
import argparse
import os
import sys
import torch.utils.data
from torch.backends import cudnn
from torch.autograd import Variable
import models
import DRO
from utils import FastRandomIdentitySampler, mkdir_if_missing, logging, display
from utils.serialization import save_checkpoint, load_checkpoint
from trainer import train
from utils import orth_reg
import DataSet
import numpy as np
import os.path as osp
cudnn.benchmark = True
use_gpu = True
# Batch Norm Freezer : bring 2% improvement on CUB
def set_bn_eval(m):
classname = m.__class__.__name__
if classname.find('BatchNorm') != -1:
m.eval()
def main(args):
print(args.p_lambda)
save_dir = args.save_dir
mkdir_if_missing(save_dir)
print("DRO:", args.DRO)
# sys.stdout: output from console
# sys.stderr: exceptions from python
sys.stdout = logging.Logger(os.path.join(save_dir, 'log.txt')) #sys.stdout --> 'log.txt'
sys.stderr = logging.Logger(os.path.join(save_dir, 'error.txt')) #sys.stderr --> 'error.txt'
display(args)
start = 0
model = models.create(args.net, pretrained=True, dim=args.dim)
save_checkpoint({
'state_dict': model.state_dict(),
'epoch': 0,
}, is_best=False, fpath=osp.join(args.save_dir, 'ckp_ep'+ str(start) + '.pth.tar'))
# for vgg and densenet
if args.resume is None:
model_dict = model.state_dict()
else:
# resume model
print('load model from {}'.format(args.resume))
chk_pt = load_checkpoint(args.resume)
weight = chk_pt['state_dict']
start = chk_pt['epoch']
model.load_state_dict(weight)
model = torch.nn.DataParallel(model)
model = model.cuda()
# freeze BN
if args.freeze_BN is True:
print(40 * '#', '\n BatchNorm frozen')
model.apply(set_bn_eval) # m represents default layers.
else:
print(40*'#', 'BatchNorm NOT frozen')
optimizer = torch.optim.Adam(model.module.parameters(), lr=args.lr,
weight_decay=args.weight_decay)
print("--------------------------:", args.p_lambda)
criterion = DRO.create(args.DRO, loss = args.loss, margin=args.margin, alpha=args.alpha,
beta = args.beta,
p_lambda = args.p_lambda, p_lambda_neg = args.p_lambda_neg, K = args.K,
select_TOPK_all = args.select_TOPK_all, p_choice = args.p_choice,
truncate_p = args.truncate_p).cuda()
# Decor_loss = losses.create('decode').cuda()
print("Train, RAE:", args.mode)
data = DataSet.create(args.data, ratio=args.ratio, width=args.width, origin_width=args.origin_width, root=args.data_root, RAE=args.mode)
train_loader = torch.utils.data.DataLoader(
data.train, batch_size=args.batch_size,
sampler=FastRandomIdentitySampler(data.train, num_instances=args.num_instances),
drop_last=True, pin_memory=True, num_workers=args.nThreads)
# save the train information
for epoch in range(start, args.epochs):
train(epoch=epoch, model=model, criterion=criterion,
optimizer=optimizer, train_loader=train_loader, args=args)
if epoch == 1:
optimizer.param_groups[0]['lr_mul'] = 0.1
if (epoch+1) % args.save_step == 0 or epoch==0:
if use_gpu:
state_dict = model.module.state_dict()
else:
state_dict = model.state_dict()
save_checkpoint({
'state_dict': state_dict,
'epoch': (epoch+1),
}, is_best=False, fpath=osp.join(args.save_dir, 'ckp_ep' + str(epoch + 1) + '.pth.tar'))
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='Deep Metric Learning')
# hype-parameters
parser.add_argument('--lr', type=float, default=1e-5, help="learning rate of new parameters")
parser.add_argument('--batch_size', '-b', default=128, type=int, metavar='N',
help='mini-batch size (1 = pure stochastic) Default: 256')
parser.add_argument('--num_instances', default=8, type=int, metavar='n',
help=' number of samples from one class in mini-batch')
parser.add_argument('--dim', default=512, type=int, metavar='n',
help='dimension of embedding space')
parser.add_argument('--width', default=224, type=int,
help='width of input image')
parser.add_argument('--origin_width', default=256, type=int,
help='size of origin image')
parser.add_argument('--ratio', default=0.16, type=float,
help='random crop ratio for train data')
parser.add_argument('--DRO', default='DRO_TOPK', type=str,help = 'which kinds of methods that we choose to use.' )
parser.add_argument('--alpha', default=30, type=float, metavar='n',
help='hyper parameter in NCA and its variants')
parser.add_argument('--beta', default=0.1, type=float, metavar='n',
help='hyper parameter in some deep metric loss functions')
parser.add_argument('--orth_reg', default=0, type=float,
help='hyper parameter coefficient for orth-reg loss')
parser.add_argument('-K', default=16, type=int, metavar='n',
help='number of neighbour points in KNN')
parser.add_argument('--margin', default=0.5, type=float,
help='margin in loss function')
parser.add_argument('--init', default='random',
help='the initialization way of FC layer')
# network
parser.add_argument('--freeze_BN', default=True, type=bool, required=False, metavar='N',
help='Freeze BN if True')
parser.add_argument('--data', default='cub', required=True,
help='name of Data Set')
parser.add_argument('--data_root', type=str, default=None,
help='path to Data Set')
parser.add_argument('--net', default='VGG16-BN')
parser.add_argument('--loss', default='branch', required=True,
help='loss for training network')
parser.add_argument('--epochs', default=600, type=int, metavar='N',
help='epochs for training process')
parser.add_argument('--save_step', default=50, type=int, metavar='N',
help='number of epochs to save model')
# Resume from checkpoint
parser.add_argument('--resume', '-r', default=None,
help='the path of the pre-trained model')
parser.add_argument('--print_freq', default=20, type=int,
help='display frequency of training')
parser.add_argument('--save_dir', default=None,
help='where the trained models save')
parser.add_argument('--nThreads', '-j', default=16, type=int, metavar='N',
help='number of data loading threads (default: 2)')
parser.add_argument('--momentum', type=float, default=0.9)
parser.add_argument('--weight-decay', type=float, default=5e-4)
parser.add_argument('--loss_base', type=float, default=0.75)
#DRO hype-parameters
parser.add_argument('--p_choice', default=1, type = int, help = 'Grouping Methods of P')
parser.add_argument('--p_lambda', default=0.1, type=float, help = 'lambda of DRO regularizer for all pairs or positive pairs')
parser.add_argument('--p_lambda_neg', default=1.0 , type =float, help = 'negative DRO regularizer')
parser.add_argument('--truncate_p', default=0, type = int, help = 'weather to truncate p')
parser.add_argument('--K', default=50, type = int, help = 'K of Top K. We are actually select 2K unique samples from the batch.')
parser.add_argument('--select_TOPK_all', default=1, type = int, help = '1: TOPK over batch. 2: TOPK by class.')
parser.add_argument('--plambda_eq', default=1, type=int, help="whether the lambda for positive pairs is equal to the lambda for negative pairs")
parser.add_argument('--mode', default='None', type=str, help="Data Augmentation Type")
main(parser.parse_args())