def _train_prepare(self):
if self.update_conv_layers != 0:
name_ls = [
x for x, p in self.model.named_parameters()
if ('fc' not in x and 'full' not in x)
]
old_name_ls = [x for x in name_ls if 'added' not in x]
update_index = 5 * self.update_conv_layers if self.update_conv_layers >= 1 else int(
len(old_name_ls) * (1 - self.update_conv_layers))
if 'r-a-ft' in self.prefix:
update_index = name_ls.index('base.layer4.1.bn3.bias')
freeze_model(self.model, update_prefix=name_ls[update_index + 1:])
self.num_train_pids, self.num_train_imgs = statistic_manager(
self.train_root)
self.num_test_pids, self.num_test_imgs = statistic_manager(
self.test_root)
self.num_test_unseen_pids, self.num_test_unseen_imgs = statistic_manager(
self.test_unseen_root)
print('\n', 'TRAINING METHOD:', method)
print(' TRAINING CLASSES:', ('%4d' % self.num_train_pids),
', COUNT: ', self.num_train_imgs)
print(' TEST CLASSES:', ('%4d' % self.num_test_pids),
', COUNT: ', self.num_test_imgs)
print('UNSEEN TEST CLASSES:', ('%4d' % self.num_test_unseen_pids),
', COUNT: ', self.num_test_unseen_imgs, '\n')
"""
Hyper-Parameters for model
"""
SEED, self.margin, self.lr, self.weight_decay, self.num_gpu, self.step_size = 0, 20.0, .1, 5e-4, 1, 40
self.save_dir = os.path.join('model', self.prefix)
if not os.path.exists(self.save_dir):
os.makedirs(self.save_dir)
with open(os.path.join(self.save_dir, 'readme.txt'), 'wb+') as f:
time_str = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
c = 'This folder is generated on %s, created by %s, V10 means 1w pics/epoch,\r\n'\
'created for saving trained model, usually,\r\n'\
'the best model should be the one named with best in it.\r\n\r\n' % (time_str, self.prefix)
f.write(c.encode())
torch.manual_seed(SEED)
self.use_gpu = torch.cuda.is_available()
if self.use_gpu:
print('Currently using GPU')
cudnn.benchmark = True
torch.cuda.manual_seed_all(SEED)
self.model = nn.DataParallel(
self.model).cuda() # use multi GPUs togather
else:
print('Currently using CPU')
self.tri_criterion = TripletLoss(self.margin)
self.optimizer = torch.optim.Adam(self.optim_policy,
lr=self.lr,
weight_decay=self.weight_decay)
def train(args):
logging.info('======= user config ======')
logging.info(pprint(opt))
logging.info(pprint(args))
logging.info('======= end ======')
train_data, test_data, num_query = get_data_provider(opt)
net = getattr(network, opt.network.name)(opt.dataset.num_classes,
opt.network.last_stride)
net = nn.DataParallel(net).cuda()
optimizer = getattr(torch.optim,
opt.train.optimizer)(net.parameters(),
lr=opt.train.lr,
weight_decay=opt.train.wd)
ce_loss = nn.CrossEntropyLoss()
triplet_loss = TripletLoss(margin=opt.train.margin)
def ce_loss_func(scores, feat, labels):
ce = ce_loss(scores, labels)
return ce
def tri_loss_func(scores, feat, labels):
tri = triplet_loss(feat, labels)[0]
return tri
def ce_tri_loss_func(scores, feat, labels):
ce = ce_loss(scores, labels)
triplet = triplet_loss(feat, labels)[0]
return ce + triplet
if opt.train.loss_fn == 'softmax':
loss_fn = ce_loss_func
elif opt.train.loss_fn == 'triplet':
loss_fn = tri_loss_func
elif opt.train.loss_fn == 'softmax_triplet':
loss_fn = ce_tri_loss_func
else:
raise ValueError('Unknown loss func {}'.format(opt.train.loss_fn))
lr_scheduler = LRScheduler(base_lr=opt.train.lr,
step=opt.train.step,
factor=opt.train.factor,
warmup_epoch=opt.train.warmup_epoch,
warmup_begin_lr=opt.train.warmup_begin_lr)
mod = Solver(opt, net)
mod.fit(train_data=train_data,
test_data=test_data,
num_query=num_query,
optimizer=optimizer,
criterion=loss_fn,
lr_scheduler=lr_scheduler)
def get_loss():
xent_criterion = nn.CrossEntropyLoss()
triplet = TripletLoss(opt.margin)
if opt.loss == 'softmax':
def criterion(feat, score, labels):
return xent_criterion(score, labels)
elif opt.loss == 'triplet':
def criterion(feat, score, labels):
return triplet(feat, labels)[0]
else:
def criterion(feat, score, labels):
return xent_criterion(score, labels) + triplet(feat, labels)[0]
return criterion
def train(**kwargs):
opt._parse(kwargs)
opt.model_name = 'bfe_test'
# set random seed and cudnn benchmark
torch.manual_seed(opt.seed)
os.makedirs(opt.save_dir, exist_ok=True)
use_gpu = torch.cuda.is_available()
sys.stdout = Logger(osp.join(opt.save_dir, 'log_train.txt'))
print('=========user config==========')
pprint(opt._state_dict())
print('============end===============')
if use_gpu:
print('currently using GPU')
cudnn.benchmark = True
torch.cuda.manual_seed_all(opt.seed)
else:
print('currently using cpu')
print('initializing dataset {}'.format(opt.dataset))
dataset = data_manager.init_dataset(name=opt.dataset, mode=opt.mode)
pin_memory = True if use_gpu else False
summary_writer = SummaryWriter(osp.join(opt.save_dir, 'tensorboard_log'))
trainloader = DataLoader(ImageData(dataset.train,
TrainTransform(opt.datatype)),
sampler=RandomIdentitySampler(
dataset.train, opt.num_instances),
batch_size=opt.train_batch,
num_workers=opt.workers,
pin_memory=pin_memory,
drop_last=True)
queryloader = DataLoader(ImageData(dataset.query,
TestTransform(opt.datatype)),
batch_size=opt.test_batch,
num_workers=opt.workers,
pin_memory=pin_memory)
galleryloader = DataLoader(ImageData(dataset.gallery,
TestTransform(opt.datatype)),
batch_size=opt.test_batch,
num_workers=opt.workers,
pin_memory=pin_memory)
queryFliploader = DataLoader(ImageData(dataset.query,
TestTransform(opt.datatype, True)),
batch_size=opt.test_batch,
num_workers=opt.workers,
pin_memory=pin_memory)
galleryFliploader = DataLoader(ImageData(dataset.gallery,
TestTransform(opt.datatype,
True)),
batch_size=opt.test_batch,
num_workers=opt.workers,
pin_memory=pin_memory)
print('initializing model ...')
model = BFE(dataset.num_train_pids, 1.0, 0.33)
optim_policy = model.get_optim_policy()
if opt.pretrained_model:
state_dict = torch.load(opt.pretrained_model)['state_dict']
# state_dict = {k: v for k, v in state_dict.items() \
# if not ('reduction' in k or 'softmax' in k)}
model.load_state_dict(state_dict, False)
print('load pretrained model ' + opt.pretrained_model)
print('model size: {:.5f}M'.format(
sum(p.numel() for p in model.parameters()) / 1e6))
if use_gpu:
model = nn.DataParallel(model).cuda()
reid_evaluator = ResNetEvaluator(model)
if opt.evaluate:
reid_evaluator.evaluate(queryloader,
galleryloader,
queryFliploader,
galleryFliploader,
re_ranking=opt.re_ranking,
savefig=opt.savefig)
return
# xent_criterion = nn.CrossEntropyLoss()
xent_criterion = CrossEntropyLabelSmooth(dataset.num_train_pids)
if opt.loss == 'triplet':
embedding_criterion = TripletLoss(opt.margin)
elif opt.loss == 'lifted':
embedding_criterion = LiftedStructureLoss(hard_mining=True)
elif opt.loss == 'weight':
embedding_criterion = Margin()
def criterion(triplet_y, softmax_y, labels):
losses = [embedding_criterion(output, labels)[0] for output in triplet_y] + \
[xent_criterion(output, labels) for output in softmax_y]
loss = sum(losses)
return loss
# get optimizer
if opt.optim == "sgd":
optimizer = torch.optim.SGD(optim_policy,
lr=opt.lr,
momentum=0.9,
weight_decay=opt.weight_decay)
else:
optimizer = torch.optim.Adam(optim_policy,
lr=opt.lr,
weight_decay=opt.weight_decay)
start_epoch = opt.start_epoch
# get trainer and evaluator
reid_trainer = cls_tripletTrainer(opt, model, optimizer, criterion,
summary_writer)
def adjust_lr(optimizer, ep):
if ep < 10:
lr = opt.lr * 0.1 * (ep / 10.0) # warm_up
elif ep < 50:
lr = opt.lr * (ep // 5 + 1)
elif ep < 200:
lr = opt.lr * 10.0
elif ep < 300:
lr = opt.lr
else:
lr = opt.lr * 0.1
for p in optimizer.param_groups:
p['lr'] = lr
# start training
best_rank1 = opt.best_rank
best_epoch = 0
for epoch in range(start_epoch, opt.max_epoch):
if opt.adjust_lr:
adjust_lr(optimizer, epoch + 1)
reid_trainer.train(epoch, trainloader)
# skip if not save model
if opt.eval_step > 0 and (epoch + 1) % opt.eval_step == 0 or (
epoch + 1) == opt.max_epoch:
if opt.mode == 'class':
rank1 = test(model, queryloader)
else:
rank1 = reid_evaluator.evaluate(queryloader, galleryloader,
queryFliploader,
galleryFliploader)
is_best = rank1 > best_rank1
if is_best:
best_rank1 = rank1
best_epoch = epoch + 1
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=is_best,
save_dir=opt.save_dir,
filename='checkpoint_ep' + str(epoch + 1) +
'.pth.tar')
print('Best rank-1 {:.1%}, achived at epoch {}'.format(
best_rank1, best_epoch))
def train(**kwargs):
opt._parse(kwargs)
# set random seed and cudnn benchmark
torch.manual_seed(opt.seed)
use_gpu = torch.cuda.is_available()
sys.stdout = Logger(osp.join(opt.save_dir, 'log_train.txt'))
print('=========user config==========')
pprint(opt._state_dict())
print('============end===============')
if use_gpu:
print('currently using GPU')
cudnn.benchmark = True
torch.cuda.manual_seed_all(opt.seed)
else:
print('currently using cpu')
print('initializing dataset {}'.format(opt.dataset))
dataset = data_manager.init_dataset(name=opt.dataset)
pin_memory = True if use_gpu else False
summary_writer = SummaryWriter(osp.join(opt.save_dir, 'tensorboard_log'))
if 'triplet' in opt.model_name:
trainloader = DataLoader(
ImageData(dataset.train, TrainTransform(opt.height, opt.width)),
sampler=RandomIdentitySampler(dataset.train, opt.num_instances),
batch_size=opt.train_batch,
num_workers=opt.workers,
pin_memory=pin_memory,
drop_last=True)
else:
trainloader = DataLoader(ImageData(
dataset.train, TrainTransform(opt.height, opt.width)),
batch_size=opt.train_batch,
shuffle=True,
num_workers=opt.workers,
pin_memory=pin_memory)
queryloader = DataLoader(ImageData(dataset.query,
TestTransform(opt.height, opt.width)),
batch_size=opt.test_batch,
num_workers=opt.workers,
pin_memory=pin_memory)
galleryloader = DataLoader(ImageData(dataset.gallery,
TestTransform(opt.height, opt.width)),
batch_size=opt.test_batch,
num_workers=opt.workers,
pin_memory=pin_memory)
print('initializing model ...')
if opt.model_name == 'softmax' or opt.model_name == 'softmax_triplet':
model, optim_policy = get_baseline_model(dataset.num_train_pids)
elif opt.model_name == 'triplet':
model, optim_policy = get_baseline_model(num_classes=None)
print('model size: {:.5f}M'.format(
sum(p.numel() for p in model.parameters()) / 1e6))
# xent_criterion = nn.CrossEntropyLoss()
xent_criterion = CrossEntropyLabelSmooth(dataset.num_train_pids)
tri_criterion = TripletLoss(opt.margin)
def cls_criterion(cls_scores, targets):
cls_loss = xent_criterion(cls_scores, targets)
return cls_loss
def triplet_criterion(feat, targets):
triplet_loss, _, _ = tri_criterion(feat, targets)
return triplet_loss
def cls_tri_criterion(cls_scores, feat, targets):
cls_loss = xent_criterion(cls_scores, targets)
triplet_loss, _, _ = tri_criterion(feat, targets)
loss = cls_loss + triplet_loss
return loss
# get optimizer
optimizer = torch.optim.Adam(optim_policy,
lr=opt.lr,
weight_decay=opt.weight_decay)
def adjust_lr(optimizer, ep):
if ep < 20:
lr = 1e-4 * (ep + 1) / 2
elif ep < 80:
lr = 1e-3 * opt.num_gpu
elif ep < 180:
lr = 1e-4 * opt.num_gpu
elif ep < 300:
lr = 1e-5 * opt.num_gpu
elif ep < 320:
lr = 1e-5 * 0.1**((ep - 320) / 80) * opt.num_gpu
elif ep < 400:
lr = 1e-6
elif ep < 480:
lr = 1e-4 * opt.num_gpu
else:
lr = 1e-5 * opt.num_gpu
for p in optimizer.param_groups:
p['lr'] = lr
start_epoch = opt.start_epoch
if use_gpu:
model = nn.DataParallel(model).cuda()
# get trainer and evaluator
if opt.model_name == 'softmax':
reid_trainer = clsTrainer(opt, model, optimizer, cls_criterion,
summary_writer)
elif opt.model_name == 'softmax_triplet':
reid_trainer = cls_tripletTrainer(opt, model, optimizer,
cls_tri_criterion, summary_writer)
elif opt.model_name == 'triplet':
reid_trainer = tripletTrainer(opt, model, optimizer, triplet_criterion,
summary_writer)
reid_evaluator = ResNetEvaluator(model)
# start training
best_rank1 = -np.inf
best_epoch = 0
for epoch in range(start_epoch, opt.max_epoch):
if opt.step_size > 0:
adjust_lr(optimizer, epoch + 1)
reid_trainer.train(epoch, trainloader)
# skip if not save model
if opt.eval_step > 0 and (epoch + 1) % opt.eval_step == 0 or (
epoch + 1) == opt.max_epoch:
rank1 = reid_evaluator.evaluate(queryloader, galleryloader)
is_best = rank1 > best_rank1
if is_best:
best_rank1 = rank1
best_epoch = epoch + 1
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=is_best,
save_dir=opt.save_dir,
filename='checkpoint_ep' + str(epoch + 1) +
'.pth.tar')
print('Best rank-1 {:.1%}, achived at epoch {}'.format(
best_rank1, best_epoch))
def main():
logger = logging.getLogger('global')
global criterion_xent, criterion_triplet, criterion_center
if os.path.exists(cfg.TRAIN.LOG_DIR):
shutil.rmtree(cfg.TRAIN.LOG_DIR)
os.makedirs(cfg.TRAIN.LOG_DIR)
init_log('global', logging.INFO) # log
add_file_handler('global', os.path.join(cfg.TRAIN.LOG_DIR, 'logs.txt'),
logging.INFO)
summary_writer = SummaryWriter(cfg.TRAIN.LOG_DIR) # visualise
dataset, train_loader, _, _ = build_data_loader()
model = BagReID_RESNET(dataset.num_train_bags)
criterion_xent = CrossEntropyLabelSmooth(dataset.num_train_bags,
use_gpu=cfg.CUDA)
criterion_triplet = TripletLoss(margin=cfg.TRAIN.MARGIN)
criterion_center = CenterLoss(dataset.num_train_bags,
cfg.MODEL.GLOBAL_FEATS +
cfg.MODEL.PART_FEATS,
use_gpu=cfg.CUDA)
if cfg.TRAIN.OPTIM == "sgd":
optimizer = torch.optim.SGD(model.parameters(),
lr=cfg.SOLVER.LEARNING_RATE,
momentum=cfg.SOLVER.MOMENTUM,
weight_decay=cfg.SOLVER.WEIGHT_DECAY)
else:
optimizer = torch.optim.Adam(model.parameters(),
lr=cfg.SOLVER.LEARNING_RATE,
weight_decay=cfg.SOLVER.WEIGHT_DECAY)
center_optimizer = torch.optim.SGD(criterion_center.parameters(),
lr=cfg.SOLVER.LEARNING_RATE_CENTER)
optimizers = [optimizer, center_optimizer]
schedulers = build_lr_schedulers(optimizers)
if cfg.CUDA:
model.cuda()
if torch.cuda.device_count() > 1:
model = torch.nn.DataParallel(model, device_ids=cfg.DEVICES)
logger.info("model prepare done")
# start training
for epoch in range(cfg.TRAIN.NUM_EPOCHS):
train(epoch, train_loader, model, criterion, optimizers,
summary_writer)
for scheduler in schedulers:
scheduler.step()
# skip if not save model
if cfg.TRAIN.EVAL_STEP > 0 and (epoch + 1) % cfg.TRAIN.EVAL_STEP == 0 \
or (epoch + 1) == cfg.TRAIN.NUM_EPOCHS:
if cfg.CUDA and torch.cuda.device_count() > 1:
state_dict = model.module.state_dict()
else:
state_dict = model.state_dict()
save_checkpoint({
'state_dict': state_dict,
'epoch': epoch + 1
},
is_best=False,
save_dir=cfg.TRAIN.SNAPSHOT_DIR,
filename='checkpoint_ep' + str(epoch + 1) + '.pth')
def _train_prepare(self):
if self.update_conv_layers != 0:
name_ls = [
x for x, p in self.model.named_parameters()
if ('fc' not in x and 'full' not in x)
]
old_name_ls = [x for x in name_ls if 'added' not in x]
update_index = 5 * self.update_conv_layers if self.update_conv_layers >= 1 else int(
len(old_name_ls) * (1 - self.update_conv_layers))
if 'r-a-ft' in self.prefix and self.model_type == 'inception3':
update_index = name_ls.index('base.layer4.1.bn3.bias')
freeze_model(self.model, update_prefix=name_ls[update_index + 1:])
self.num_train_pids, self.num_train_imgs = statistic_manager(
self.train_root)
self.num_test_pids, self.num_test_imgs = statistic_manager(
self.test_root)
self.num_test_unseen_pids, self.num_test_unseen_imgs = statistic_manager(
self.test_unseen_root)
print('\n')
print(' TRAINING METHOD:', self.method)
print(' TRAINING CLASSES:', ('%4d' % self.num_train_pids),
', COUNT: ', self.num_train_imgs)
print(' TEST CLASSES:', ('%4d' % self.num_test_pids),
', COUNT: ', self.num_test_imgs)
print('UNSEEN TEST CLASSES:', ('%4d' % self.num_test_unseen_pids),
', COUNT: ', self.num_test_unseen_imgs, '\n')
"""
Hyper-Parameters for model
"""
SEED, self.margin, self.lr, self.weight_decay, self.num_gpu, self.step_size = 0, 20.0, .1, 5e-4, 1, 40
self.save_dir = os.path.join('model', self.prefix)
if not os.path.exists(self.save_dir):
os.makedirs(self.save_dir)
with open(os.path.join(self.save_dir, 'readme.txt'), 'ab+') as f:
time_str = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
c = 'This folder is generated on {}, marked by {},\r\n'\
'Here are some settings:\r\n'\
'train method: {}\r\n'\
'neural network: {}\r\n'\
'batch_size: {}\r\n'\
'training set: {}({} classes, {} images)\r\n'\
'test set: {}({} classes, {} images)\r\n'\
'input size: {}x{}\r\n'\
'and i load a model: {}\r\n\r\n'\
'And, this folder is created for saving trained model, usually,\r\n'\
'the best model should be the one named with "model.tar" in it.\r\n\r\n\r\n' .format (time_str, self.prefix,
self.method,
self.model_type,
self.batch_size,
self.train_root[21:], self.num_train_pids, self.num_train_imgs,
self.test_root[21:], self.num_test_pids, self.num_test_imgs,
self.w, self.h,
str(self.load_model_path))
f.write(c.encode())
torch.manual_seed(SEED)
self.use_gpu = torch.cuda.is_available()
if self.use_gpu:
print('Currently using GPU')
cudnn.benchmark = True
torch.cuda.manual_seed_all(SEED)
self.model = nn.DataParallel(
self.model).cuda() # use multi GPUs togather
else:
print('Currently using CPU')
self.tri_criterion = TripletLoss(self.margin)
self.optimizer = torch.optim.Adam(self.optim_policy,
lr=self.lr,
weight_decay=self.weight_decay)
def train(**kwargs):
opt._parse(kwargs)
# torch.backends.cudnn.deterministic = True # I think this line may slow down the training process
# set random seed and cudnn benchmark
torch.manual_seed(opt.seed)
random.seed(opt.seed)
np.random.seed(opt.seed)
use_gpu = torch.cuda.is_available()
sys.stdout = Logger(
os.path.join('./pytorch-ckpt/current', opt.save_dir, 'log_train.txt'))
if use_gpu:
print('currently using GPU')
cudnn.benchmark = True
else:
print('currently using cpu')
print(opt._state_dict())
print('initializing dataset {}'.format(opt.trainset_name))
if opt.trainset_name == 'combine':
#input dataset name as 'datasets'
train_dataset = data_manager.init_combine_dataset(
name=opt.trainset_name,
options=opt,
datasets=opt.datasets,
num_bn_sample=opt.batch_num_bn_estimatation * opt.test_batch,
share_cam=opt.share_cam,
num_pids=opt.num_pids)
elif opt.trainset_name == 'unreal':
# input dataset dir in 'datasets'
train_dataset = data_manager.init_unreal_dataset(
name=opt.trainset_name,
datasets=opt.datasets,
num_pids=opt.num_pids,
num_cams=opt.num_cams,
img_per_person=opt.img_per_person)
else:
train_dataset = data_manager.init_dataset(
name=opt.trainset_name,
num_bn_sample=opt.batch_num_bn_estimatation * opt.test_batch,
num_pids=opt.num_pids)
pin_memory = True if use_gpu else False
summary_writer = SummaryWriter(
os.path.join('./pytorch-ckpt/current', opt.save_dir,
'tensorboard_log'))
if opt.cam_bal:
IDSampler = IdentityCameraSampler
else:
IDSampler = IdentitySampler
if opt.trainset_name == 'combine':
samp = IDSampler(train_dataset.train, opt.train_batch,
opt.num_instances, train_dataset.cams_of_dataset,
train_dataset.len_of_real_dataset)
else:
samp = IDSampler(train_dataset.train, opt.train_batch,
opt.num_instances)
trainloader = DataLoader(data_manager.init_datafolder(
opt.trainset_name, train_dataset.train,
TrainTransform(opt.height, opt.width)),
sampler=samp,
batch_size=opt.train_batch,
num_workers=opt.workers,
pin_memory=pin_memory,
drop_last=True,
collate_fn=NormalCollateFn())
print('initializing model ...')
num_pid = train_dataset.num_train_pids if opt.loss == 'softmax' else None
model = ResNetBuilder(num_pid)
if opt.model_path is not None and 'moco' in opt.model_path:
model = load_moco_model(model, opt.model_path)
elif opt.model_path is not None:
model = load_previous_model(model,
opt.model_path,
load_fc_layers=False)
optim_policy = model.get_optim_policy()
print('model size: {:.5f}M'.format(
sum(p.numel() for p in model.parameters()) / 1e6))
if use_gpu:
model = CamDataParallel(model).cuda()
xent = nn.CrossEntropyLoss()
triplet = TripletLoss()
def standard_cls_criterion(feat, preditions, targets, global_step,
summary_writer):
identity_loss = xent(preditions, targets)
identity_accuracy = torch.mean(
(torch.argmax(preditions, dim=1) == targets).float())
summary_writer.add_scalar('cls_loss', identity_loss.item(),
global_step)
summary_writer.add_scalar('cls_accuracy', identity_accuracy.item(),
global_step)
return identity_loss
def triplet_criterion(feat, preditons, targets, global_step,
summary_writer):
triplet_loss, acc = triplet(feat, targets)
summary_writer.add_scalar('loss', triplet_loss.item(), global_step)
print(np.mean(acc.item()))
summary_writer.add_scalar('accuracy', acc.item(), global_step)
return triplet_loss
# get trainer and evaluator
optimizer, adjust_lr = get_our_optimizer_strategy(opt, optim_policy)
if opt.loss == 'softmax':
crit = standard_cls_criterion
elif opt.loss == 'triplet':
crit = triplet_criterion
reid_trainer = CameraClsTrainer(opt, model, optimizer, crit,
summary_writer)
print('Start training')
for epoch in range(opt.max_epoch):
adjust_lr(optimizer, epoch)
reid_trainer.train(epoch, trainloader)
if (epoch + 1) % opt.save_step == 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,
},
save_dir=os.path.join('./pytorch-ckpt/current',
opt.save_dir),
ep=epoch + 1)
# if (epoch+1)%15==0:
# save_checkpoint({
# 'state_dict': state_dict,
# 'epoch': epoch + 1,
# }, save_dir=os.path.join('./pytorch-ckpt/current', opt.save_dir))
if use_gpu:
state_dict = model.module.state_dict()
else:
state_dict = model.state_dict()
save_checkpoint({
'state_dict': state_dict,
'epoch': epoch + 1,
},
save_dir=os.path.join('./pytorch-ckpt/current',
opt.save_dir))
if use_gpu:
model = nn.DataParallel(model).cuda()
reid_evaluator = ResNetEvaluator(model)
if opt.evaluate:
reid_evaluator.evaluate(queryloader, galleryloader,
queryFliploader, galleryFliploader, re_ranking=opt.re_ranking, savefig=opt.savefig)
return
#xent_criterion = nn.CrossEntropyLoss()
#print(dataset.num_train_pids)
xent_criterion = CrossEntropyLabelSmooth(dataset.num_train_pids)
xent_criterion2 = CrossEntropyLabelSmooth2(dataset.num_train_pids)
if opt.loss == 'triplet':
embedding_criterion = TripletLoss(opt.margin)
elif opt.loss == 'lifted':
embedding_criterion = LiftedStructureLoss(hard_mining=True)
elif opt.loss == 'weight':
embedding_criterion = Margin()
def criterion(a ,b, c, d, e, f, triplet_y, softmax_y, labels):
losses = [embedding_criterion(output, labels)[0] for output in triplet_y] + \
[xent_criterion(output, labels) for output in softmax_y]+\
[xent_criterion2(i, j, labels) for i, j in zip(a,a)]+\
[xent_criterion2(i, j, labels) for i, j in zip(b,b)]+\
[xent_criterion2(i, j, labels) for i, j in zip(c,c)]+\
[xent_criterion2(i, j, labels) for i, j in zip(d,d)]+\
[xent_criterion2(i, j, labels) for i, j in zip(e,e)]+\
[xent_criterion2(i, j, labels) for i, j in zip(f,f)]
loss = sum(losses)
def trainer(data_pth, a, b, _time=0, layers=18):
seed = 0
# dataset options
height, width = 128, 128
# optimization options
optim = 'Adam'
max_epoch = 20
train_batch = 64
test_batch = 64
lr = 0.1
step_size = 40
gamma = 0.1
weight_decay = 5e-4
momentum = 0.9
test_margin = b
margin = a
num_instances = 4
num_gpu = 1
# model options
last_stride = 1
pretrained_model_18 = 'model/resnet18-5c106cde.pth'
pretrained_model_50 = 'model/resnet50-19c8e357.pth'
pretrained_model_34 = 'model/resnet34-333f7ec4.pth'
pretrained_model_101 = 'model/resnet101-5d3b4d8f.pth'
pretrained_model_152 = 'model/resnet152-b121ed2d.pth'
# miscs
print_freq = 10
eval_step = 1
save_dir = 'model/pytorch-ckpt/time%d' % _time
workers = 1
torch.manual_seed(seed)
use_gpu = torch.cuda.is_available()
if use_gpu:
print('currently using GPU')
cudnn.benchmark = True
torch.cuda.manual_seed_all(seed)
else:
print('currently using cpu')
pin_memory = True if use_gpu else False
# model, optim_policy = get_baseline_model(model_path=pretrained_model)
if layers == 18:
model, optim_policy = get_baseline_model(
model_path=pretrained_model_18, layers=18)
else:
model, optim_policy = get_baseline_model(
model_path=pretrained_model_50, layers=50)
# model, optim_policy = get_baseline_model(model_path=pretrained_model_18, layers=18)
# model, optim_policy = get_baseline_model(model_path=pretrained_model_34, layers=34)
# model, optim_policy = get_baseline_model(model_path=pretrained_model_101, layers=101)
# model = load_model(model, model_path='./model/pytorch-ckpt/87_layers18_margin20_epoch87.tar')
print('model\'s parameters size: {:.5f} M'.format(
sum(p.numel() for p in model.parameters()) / 1e6))
tri_criterion = TripletLoss(margin)
# get optimizer
optimizer = torch.optim.Adam(optim_policy,
lr=lr,
weight_decay=weight_decay)
def adjust_lr(optimizer, ep):
if ep < 20:
lr = 1e-4 * (ep + 1) / 2
elif ep < 80:
lr = 1e-3 * num_gpu
elif ep < 180:
lr = 1e-4 * num_gpu
elif ep < 300:
lr = 1e-5 * num_gpu
elif ep < 320:
lr = 1e-5 * 0.1**((ep - 320) / 80) * num_gpu
elif ep < 400:
lr = 1e-6
elif ep < 480:
lr = 1e-4 * num_gpu
else:
lr = 1e-5 * num_gpu
for p in optimizer.param_groups:
p['lr'] = lr
if use_gpu:
model = nn.DataParallel(model).cuda()
max_acc = .0
for epoch in range(max_epoch):
if step_size > 0:
adjust_lr(optimizer, epoch + 1)
next_margin = margin
# skip if not save model
if eval_step > 0 and (epoch + 1) % eval_step == 0 or (epoch +
1) == max_epoch:
_t1 = time.time()
train(model,
optimizer,
tri_criterion,
epoch + 1,
print_freq,
None,
data_pth=data_pth)
_t2 = time.time()
print('time for training:', '%.2f' % (_t2 - _t1), 's')
acc = evaluate_model(model, margin=20, epoch=1)
if acc > max_acc:
max_acc = acc
print('max acc:', max_acc, ', epoch:', epoch + 1)
if use_gpu:
state_dict = model.module.state_dict()
else:
state_dict = model.state_dict()
save_model_name = 'layers{}_margin{}_epoch{}.tar'.format(
layers, margin, epoch + 1)
save_checkpoint(
{
'state_dict': state_dict,
'epoch': epoch + 1,
},
is_best=False,
save_dir=save_dir,
filename=save_model_name)
margin = next_margin
return save_model_name
def trainer(data_pth, a, b, _time=0, layers=18):
seed = 0
# dataset options
height = 128
width = 128
# optimization options
optim = 'Adam'
max_epoch = 20
train_batch = 64
test_batch = 64
lr = 0.1
step_size = 40
gamma = 0.1
weight_decay = 5e-4
momentum = 0.9
test_margin = b
margin = a
num_instances = 4
num_gpu = 1
# model options
last_stride = 1
pretrained_model_18 = 'model/resnet18-5c106cde.pth'
pretrained_model_50 = 'model/resnet50-19c8e357.pth'
pretrained_model_34 = 'model/resnet34-333f7ec4.pth'
pretrained_model_101 = 'model/resnet101-5d3b4d8f.pth'
pretrained_model_152 = 'model/resnet152-b121ed2d.pth'
# miscs
print_freq = 20
eval_step = 1
save_dir = 'model/pytorch-ckpt/time%d' % _time
workers = 1
start_epoch = 0
torch.manual_seed(seed)
use_gpu = torch.cuda.is_available()
if use_gpu:
print('currently using GPU')
cudnn.benchmark = True
torch.cuda.manual_seed_all(seed)
else:
print('currently using cpu')
pin_memory = True if use_gpu else False
print('initializing dataset {}'.format('Tableware'))
dataset = Tableware(data_pth)
trainloader = DataLoader(
ImageData(dataset.train, TrainTransform(height, width)),
batch_size=train_batch, num_workers=workers,
pin_memory=pin_memory, drop_last=True
)
# testloader = DataLoader(
# ImageData(dataset.test, TestTransform(height, width)),
# batch_size=test_batch, num_workers=workers,
# pin_memory=pin_memory, drop_last=True
# )
# model, optim_policy = get_baseline_model(model_path=pretrained_model)
if layers == 18:
model, optim_policy = get_baseline_model(model_path=pretrained_model_18, layers=18)
else:
model, optim_policy = get_baseline_model(model_path=pretrained_model_50, layers=50)
# model, optim_policy = get_baseline_model(model_path=pretrained_model_18, layers=18)
# model, optim_policy = get_baseline_model(model_path=pretrained_model_34, layers=34)
# model, optim_policy = get_baseline_model(model_path=pretrained_model_101, layers=101)
# model = load_model(model, model_path='./model/pytorch-ckpt/87_layers18_margin20_epoch87.tar')
print('model\'s parameters size: {:.5f} M'.format(sum(p.numel() for p in model.parameters()) / 1e6))
inner_dist = 0
outer_dist = 0
max_outer = 0
min_outer = 0
max_iner = 0
min_iner = 0
tri_criterion = TripletLoss(margin)
# get optimizer
optimizer = torch.optim.Adam(
optim_policy, lr=lr, weight_decay=weight_decay
)
def adjust_lr(optimizer, ep):
if ep < 20:
lr = 1e-4 * (ep + 1) / 2
elif ep < 80:
lr = 1e-3 * num_gpu
elif ep < 180:
lr = 1e-4 * num_gpu
elif ep < 300:
lr = 1e-5 * num_gpu
elif ep < 320:
lr = 1e-5 * 0.1 ** ((ep - 320) / 80) *num_gpu
elif ep < 400:
lr = 1e-6
elif ep < 480:
lr = 1e-4 * num_gpu
else:
lr = 1e-5 * num_gpu
for p in optimizer.param_groups:
p['lr'] = lr
if use_gpu:
model = nn.DataParallel(model).cuda()
evaluator = Evaluator(model)
for epoch in range(start_epoch, max_epoch):
if step_size > 0:
adjust_lr(optimizer, epoch + 1)
next_margin = margin
# skip if not save model
if eval_step > 0 and (epoch + 1) % eval_step == 0 or (epoch + 1) == max_epoch:
save_record_path = 'margin_'+ str(margin) + '_epoch_' + str(epoch + 1) + '.txt'
_t1 =time.time()
train(model, optimizer, tri_criterion, epoch, print_freq, trainloader, data_pth=data_pth)
_t2 = time.time()
print('time for training:', '%.2f' % (_t2 - _t1), 's')
"""
acc, inner_dist, outer_dist, max_outer, min_outer, max_iner, min_iner = evaluator.evaluate(testloader, test_margin, save_record_path)
print('margin:{}, epoch:{}, acc:{}'.format(margin, epoch+1, acc))
f = open('record.txt', 'a')
f.write('margin:{}, epoch:{}, acc:{}\n'.format(margin, epoch+1, acc))
f.close()
"""
is_best = False
# save_model_path = 'new_margin({})_epoch({}).pth.tar'.format(margin, epoch+1)
save_model_path = 'time{}_layers{}_margin{}_epoch{}.tar'.format(_time, layers, margin, epoch+1)
# save_model_path = 'layers34_margin{}_epoch{}.tar'.format(margin, epoch+1)
# save_model_path = 'layers101_margin{}_epoch{}.tar'.format(margin, epoch+1)
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=is_best, save_dir=save_dir, filename=save_model_path)
model.eval()
acc = do_get_feature_and_t(model, margin=20, epoch=1)
margin = next_margin
return save_model_path, inner_dist, outer_dist, max_outer, min_outer, max_iner, min_iner
def save_network(network, epoch_label):
save_filename = 'net_%s.pth' % epoch_label
save_path = os.path.join('./models0', save_filename)
torch.save(network.cpu().state_dict(), save_path)
if torch.cuda.is_available:
network.cuda(gpu_ids[0])
model = Model(num_classes=len(class_names))
if use_gpu:
model = model.cuda()
id_criterion = nn.CrossEntropyLoss()
tri_criterion = TripletLoss(margin=0.6)
optimizer_ft = optim.Adam(model.parameters(), lr=opt.lr, weight_decay=1e-5)
dir_name = os.path.join('./models0', '')
if not os.path.isdir(dir_name):
os.makedirs(dir_name)
with open('%s/opts.json' % dir_name, 'w') as fp:
json.dump(vars(opt), fp, indent=1)
sys.stdout = Logger(os.path.join(dir_name, 'log.txt'))
model = train_model(model, id_criterion, tri_criterion, optimizer_ft, num_epochs=300)
def train(**kwargs):
opt._parse(kwargs)
# set random seed and cudnn benchmark
sys.stdout = Logger(osp.join(opt.save_dir, 'log_train.txt'))
print('=========user config==========')
pprint(opt._state_dict())
print('============end===============')
print('initializing dataset {}'.format(opt.dataset))
dataset = data_manager.init_dataset(name=opt.dataset)
summary_writer = SummaryWriter(osp.join(opt.save_dir, 'tensorboard_log'))
if 'triplet' in opt.model_name:
trainloader = DataLoader(
ImageData(dataset.train, TrainTransform(opt.height, opt.width)),
sampler=RandomIdentitySampler(dataset.train, opt.num_instances),
batch_size=opt.train_batch,
num_workers=opt.workers,
last_batch='discard')
else:
trainloader = DataLoader(
ImageData(dataset.train, TrainTransform(opt.height, opt.width)),
batch_size=opt.train_batch,
shuffle=True,
num_workers=opt.workers,
)
queryloader = DataLoader(
ImageData(dataset.query, TestTransform(opt.height, opt.width)),
batch_size=opt.test_batch,
num_workers=opt.workers,
)
galleryloader = DataLoader(
ImageData(dataset.gallery, TestTransform(opt.height, opt.width)),
batch_size=opt.test_batch,
num_workers=opt.workers,
)
print('initializing model ...')
model = get_baseline_model(dataset.num_train_pids, mx.gpu(0),
opt.pretrained_model)
print('model size: {:.5f}M'.format(
sum(p.data().size for p in model.collect_params().values()) / 1e6))
xent_criterion = gluon.loss.SoftmaxCrossEntropyLoss()
tri_criterion = TripletLoss(opt.margin)
def cls_criterion(cls_scores, feat, targets):
cls_loss = xent_criterion(cls_scores, targets)
return cls_loss
def triplet_criterion(cls_scores, feat, targets):
triplet_loss, dist_ap, dist_an = tri_criterion(feat, targets)
return triplet_loss
def cls_tri_criterion(cls_scores, feat, targets):
cls_loss = xent_criterion(cls_scores, targets)
triplet_loss, dist_ap, dist_an = tri_criterion(feat, targets)
loss = cls_loss + triplet_loss
return loss
# get optimizer
optimizer = gluon.Trainer(model.collect_params(), opt.optim, {
'learning_rate': opt.lr,
'wd': opt.weight_decay
})
def adjust_lr(optimizer, ep):
if ep < 20:
lr = 1e-4 * (ep + 1) / 2
elif ep < 80:
lr = 1e-3 * opt.num_gpu
elif ep < 180:
lr = 1e-4 * opt.num_gpu
elif ep < 300:
lr = 1e-5 * opt.num_gpu
elif ep < 320:
lr = 1e-5 * 0.1**((ep - 320) / 80) * opt.num_gpu
elif ep < 400:
lr = 1e-6
elif ep < 480:
lr = 1e-4 * opt.num_gpu
else:
lr = 1e-5 * opt.num_gpu
optimizer.set_learning_rate(lr)
start_epoch = opt.start_epoch
# get trainer and evaluator
use_criterion = None
if opt.model_name == 'softmax':
use_criterion = cls_criterion
elif opt.model_name == 'softmax_triplet':
use_criterion = cls_tri_criterion
elif opt.model_name == 'triplet':
use_criterion = triplet_criterion
reid_trainer = reidTrainer(opt, model, optimizer, use_criterion,
summary_writer, mx.gpu(0))
reid_evaluator = reidEvaluator(model, mx.gpu(0))
# start training
best_rank1 = -np.inf
best_epoch = 0
for epoch in range(start_epoch, opt.max_epoch):
if opt.step_size > 0:
adjust_lr(optimizer, epoch + 1)
reid_trainer.train(epoch, trainloader)
# skip if not save model
if opt.eval_step > 0 and (epoch + 1) % opt.eval_step == 0 or (
epoch + 1) == opt.max_epoch:
rank1 = reid_evaluator.evaluate(queryloader, galleryloader)
is_best = rank1 > best_rank1
if is_best:
best_rank1 = rank1
best_epoch = epoch + 1
state_dict = {'model': model, 'epoch': epoch}
save_checkpoint(state_dict,
is_best=is_best,
save_dir=opt.save_dir,
filename='checkpoint_ep' + str(epoch + 1) +
'.params')
print('Best rank-1 {:.1%}, achived at epoch {}'.format(
best_rank1, best_epoch))
def trainer(data_pth):
seed = 0
# dataset options
height = 128
width = 128
# optimization options
optim = 'Adam'
max_epoch = 1
train_batch = 64
test_batch = 64
lr = 0.1
step_size = 40
gamma = 0.1
weight_decay = 5e-4
momentum = 0.9
test_margin = 10.0
margin = 1.0
num_instances = 4
num_gpu = 1
# model options
last_stride = 1
pretrained_model = 'model/resnet50-19c8e357.pth'
# miscs
print_freq = 20
eval_step = 1
save_dir = 'model/pytorch-ckpt/'
workers = 1
start_epoch = 0
torch.manual_seed(seed)
use_gpu = torch.cuda.is_available()
if use_gpu:
print('currently using GPU')
cudnn.benchmark = True
torch.cuda.manual_seed_all(seed)
else:
print('currently using cpu')
pin_memory = True if use_gpu else False
print('initializing dataset {}'.format('Tableware'))
dataset = Tableware(data_pth)
trainloader = DataLoader(ImageData(dataset.train,
TrainTransform(height, width)),
batch_size=train_batch,
num_workers=workers,
pin_memory=pin_memory,
drop_last=True)
testloader = DataLoader(ImageData(dataset.test,
TestTransform(height, width)),
batch_size=test_batch,
num_workers=workers,
pin_memory=pin_memory,
drop_last=True)
model, optim_policy = get_baseline_model(model_path=pretrained_model)
print('model size: {:.5f}M'.format(
sum(p.numel() for p in model.parameters()) / 1e6))
tri_criterion = TripletLoss(margin)
# get optimizer
optimizer = torch.optim.Adam(optim_policy,
lr=lr,
weight_decay=weight_decay)
def adjust_lr(optimizer, ep):
if ep < 20:
lr = 1e-4 * (ep + 1) / 2
elif ep < 80:
lr = 1e-3 * num_gpu
elif ep < 180:
lr = 1e-4 * num_gpu
elif ep < 300:
lr = 1e-5 * num_gpu
elif ep < 320:
lr = 1e-5 * 0.1**((ep - 320) / 80) * num_gpu
elif ep < 400:
lr = 1e-6
elif ep < 480:
lr = 1e-4 * num_gpu
else:
lr = 1e-5 * num_gpu
for p in optimizer.param_groups:
p['lr'] = lr
if use_gpu:
model = nn.DataParallel(model).cuda()
evaluator = Evaluator(model)
# start training
best_acc = -np.inf
best_epoch = 0
for epoch in range(start_epoch, max_epoch):
if step_size > 0:
adjust_lr(optimizer, epoch + 1)
train(model, optimizer, tri_criterion, epoch, print_freq, trainloader)
# skip if not save model
if eval_step > 0 and (epoch + 1) % eval_step == 0 or (epoch +
1) == max_epoch:
acc = evaluator.evaluate(testloader, test_margin)
is_best = acc > best_acc
if is_best:
best_acc = acc
best_epoch = epoch + 1
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=is_best,
save_dir=save_dir,
filename='checkpoint_ep' + str(epoch + 1) +
'.pth.tar')
print('Best accuracy {:.1%}, achieved at epoch {}'.format(
best_acc, best_epoch))
def train(**kwargs):
#### Part 1 : Initialization
opt._parse(kwargs)
torch.backends.cudnn.deterministic = True
# set random seed and cudnn benchmark
#torch.manual_seed(opt.seed)
#random.seed(opt.seed)
#np.random.seed(opt.seed)
use_gpu = torch.cuda.is_available()
sys.stdout = Logger(osp.join(opt.save_dir, 'log_train.txt'))
print('=========user config==========')
pprint(opt._state_dict())
print('============end===============')
if use_gpu:
print('currently using GPU')
cudnn.benchmark = True
torch.cuda.manual_seed_all(opt.seed)
else:
print('currently using cpu')
#### Part 2 : Preparing Data
print('initializing train dataset {}'.format(opt.trainset))
train_dataset = data_manager.init_dataset(name=opt.trainset)
print('initializing test dataset {}'.format(opt.testset))
test_dataset = data_manager.init_dataset(name=opt.testset)
pin_memory = True if use_gpu else False
pin_memory = False
summary_writer = SummaryWriter(osp.join(opt.save_dir, 'tensorboard_log'))
collateFn = NormalCollateFn()
if opt.sampler == "randomidentity":
trainloader = DataLoader(
data_manager.init_datafolder(
opt.trainset,
train_dataset.train,
TrainTransform(opt.height,
opt.width,
random_erase=opt.with_randomerase),
if_train=True),
sampler=RandomIdentitySampler(train_dataset.train,
opt.num_instances),
batch_size=opt.train_batch,
num_workers=opt.workers,
pin_memory=pin_memory,
drop_last=True,
collate_fn=collateFn,
)
elif opt.sampler == "randomidentitycamera":
trainloader = DataLoader(
data_manager.init_datafolder(
opt.trainset,
train_dataset.train,
TrainTransform(opt.height,
opt.width,
random_erase=opt.with_randomerase),
if_train=True),
batch_sampler=RandomIdentityCameraSampler(train_dataset.train,
opt.num_instances,
opt.train_batch),
num_workers=opt.workers,
pin_memory=pin_memory,
collate_fn=collateFn,
)
queryloader = DataLoader(data_manager.init_datafolder(opt.testset,
test_dataset.query,
TestTransform(
opt.height,
opt.width),
if_train=False),
batch_size=opt.test_batch,
num_workers=opt.workers,
pin_memory=pin_memory)
galleryloader = DataLoader(data_manager.init_datafolder(
opt.testset,
test_dataset.gallery,
TestTransform(opt.height, opt.width),
if_train=False),
batch_size=opt.test_batch,
num_workers=opt.workers,
pin_memory=pin_memory)
#### Part 3 : Preparing Backbone Network
print('initializing model ...')
if opt.model_name in ['triplet', 'distance']:
model, optim_policy = get_baseline_model(num_classes=None,
model='triplet')
elif opt.model_name in ["softmax"]:
model, optim_policy = get_baseline_model(train_dataset.num_train_pids,
model='softmax',
drop_prob=opt.drop)
else:
assert False, "unknown model name"
if (not opt.model_path == 'zero') and 'tar' in opt.model_path:
print('load pretrain reid model......' + opt.model_path)
ckpt = torch.load(opt.model_path)
# remove classifer
tmp = dict()
for k, v in ckpt['state_dict'].items():
if opt.keep_layer:
for i in opt.keep_layer:
if 'layer' + str(i) in k:
#print(k+" skip....")
continue
if opt.keepfc or ('fc' not in k and 'classifier' not in k):
tmp[k] = v
ckpt['state_dict'] = tmp
model.load_state_dict(ckpt['state_dict'], strict=False)
print('model size: {:.5f}M'.format(
sum(p.numel() for p in model.parameters()) / 1e6))
#### Part 4: Preparing Loss Functions
if opt.margin1 is not None:
distance_loss = DistanceLoss(margin=(opt.margin1, opt.margin2))
else:
distance_loss = DistanceLoss()
tri_loss = TripletLoss(margin=opt.margin)
xent_loss = nn.CrossEntropyLoss()
vis = dict()
vis['tri_acc1'] = AverageMeter()
vis['tri_acc2'] = AverageMeter()
vis['cls_accuracy'] = AverageMeter()
vis['cls_loss'] = AverageMeter()
def dist_criterion(feat,
targets,
cameras,
model=None,
paths=None,
epoch=0):
dis_loss, tri_acc1, tri_acc2 = distance_loss(feat,
targets,
cameras,
model,
paths,
epoch=epoch)
vis['tri_acc1'].update(float(tri_acc1))
vis['tri_acc2'].update(float(tri_acc2))
return dis_loss
def triplet_criterion(feat, targets):
triplet_loss, tri_accuracy, _, _ = tri_loss(feat, targets)
vis['tri_acc1'].update(float(tri_accuracy))
return triplet_loss
def cls_criterion(cls_scores, targets):
cls_loss = xent_loss(cls_scores, targets)
_, preds = torch.max(cls_scores.data, 1)
corrects = float(torch.sum(preds == targets.data))
vis['cls_accuracy'].update(float(corrects / opt.train_batch))
vis['cls_loss'].update(float(cls_loss))
return cls_loss
#### Part 5: Preparing Optimizer and Trainer
optimizer, adjust_lr = get_optimizer_strategy(opt.model_name, optim_policy,
opt)
start_epoch = opt.start_epoch
if use_gpu:
model = nn.DataParallel(model).cuda()
#model=model.cuda()
# get trainer and evaluatori
if opt.model_name == "distance":
reid_trainer = tripletTrainer(opt,
model,
optimizer,
dist_criterion,
summary_writer,
need_cam=True)
elif opt.model_name == 'triplet' or opt.model_name == 'triplet_fc':
reid_trainer = tripletTrainer(opt, model, optimizer, triplet_criterion,
summary_writer)
elif opt.model_name == 'softmax':
reid_trainer = clsTrainer(opt, model, optimizer, cls_criterion,
summary_writer)
else:
print("Error: Unknown model name {}".format(opt.model_name))
reid_evaluator = evaluator_manager.init_evaluator(opt.testset,
model,
flip=True)
#### Part 6 : Training
best_rank1 = -np.inf
best_epoch = 0
for epoch in range(start_epoch, opt.max_epoch):
if opt.step_size > 0:
current_lr = adjust_lr(optimizer, epoch)
reid_trainer.train(epoch, trainloader)
for k, v in vis.items():
print("{}:{}".format(k, v.mean))
v.reset()
if (epoch + 1) == opt.max_epoch:
if use_gpu and opt.num_gpu > 1:
state_dict = model.module.state_dict()
else:
state_dict = model.state_dict()
save_checkpoint({
'state_dict': state_dict,
'epoch': epoch + 1,
},
is_best=False,
save_dir=opt.save_dir,
filename='checkpoint_ep' + str(epoch + 1) +
'.pth.tar')
# skip if not save model
if (opt.eval_step > 0 and
(epoch + 1) % opt.eval_step == 0 and epoch >= 0
or (epoch + 1) == opt.max_epoch):
#print('Test on '+opt.testset)
#rank1 = reid_evaluator.evaluate(queryloader, galleryloader,normalize=opt.with_normalize)
print('Test on ' + opt.trainset)
if use_gpu and opt.num_gpu > 1:
state_dict = model.module.state_dict()
else:
state_dict = model.state_dict()
save_checkpoint({
'state_dict': state_dict,
'epoch': epoch + 1,
},
is_best=False,
save_dir=opt.save_dir,
filename='checkpoint_ep' + str(epoch + 1) +
'.pth.tar')
rank1, mAP = reid_evaluator.evaluate(queryloader,
galleryloader,
normalize=opt.with_normalize)
is_best = rank1 > best_rank1
if is_best:
best_rank1 = rank1
best_epoch = epoch + 1
save_checkpoint(
{
'state_dict': state_dict,
'epoch': epoch + 1,
},
is_best=False,
save_dir=opt.save_dir,
filename='checkpoint_ep' + str(epoch + 1) + '.pth.tar')
print('Best rank-1 {:.1%}, achieved at epoch {}'.format(
best_rank1, best_epoch))