def __init__(
self,
datamanager,
model,
optimizer,
margin=0.3,
weight_t=1,
weight_x=1,
weight_r = 0.0001,
scheduler=None,
use_gpu=True,
label_smooth=True
):
super(ImageTripletAEEngine, self
).__init__(datamanager, model, optimizer, scheduler, use_gpu)
self.weight_t = weight_t
self.weight_x = weight_x
self.weight_r = weight_r
self.criterion_t = TripletLoss(margin=margin)
self.criterion_x = CrossEntropyLoss(
num_classes=self.datamanager.num_train_pids,
use_gpu=self.use_gpu,
label_smooth=label_smooth
)
self.criterion_mse = torch.nn.MSELoss()
self.random = RandomErasing(probability=0.5)
self.mgn_loss = Loss(num_classes=self.datamanager.num_train_pids,use_gpu=self.use_gpu,label_smooth=label_smooth)
self.BCE_criterion = torch.nn.BCEWithLogitsLoss()
def __init__(self,
datamanager,
model,
optimizer,
margin=0.3,
weight_t=1,
weight_x=1,
scheduler=None,
use_gpu=True,
label_smooth=True,
conf_penalty=0.0):
super(ImageTripletEngine, self).__init__(datamanager, model, optimizer,
scheduler, use_gpu)
self.weight_t = weight_t
self.weight_x = weight_x
self.criterion_t = TripletLoss(margin=margin)
self.criterion_x = CrossEntropyLoss(use_gpu=self.use_gpu,
label_smooth=label_smooth,
conf_penalty=conf_penalty)
assert len(self.models) == 1
self.model = self.models['model']
self.optimizer = self.optims['model']
def __init__(self,
datamanager,
model,
optimizer,
margin=0.3,
weight_x=1,
weight_t=1,
weight_r=1,
weight_a=1,
scheduler=None,
use_gpu=True,
label_smooth=True,
swap_size=(8, 4)):
super(ImageSoftmaxDCLTripletEngine,
self).__init__(datamanager, use_gpu)
self.model = model
self.optimizer = optimizer
self.scheduler = scheduler
self.register_model('model', model, optimizer, scheduler)
self.weight_t = weight_t
self.weight_x = weight_x
self.weight_r = weight_r
self.weight_a = weight_a
self.criterion_x = CrossEntropyLoss(
num_classes=self.datamanager.num_train_pids,
use_gpu=self.use_gpu,
label_smooth=label_smooth)
self.criterion_trip = TripletLoss(margin=0.7)
self.criterion_rec = ReconstructionLoss()
self.criterion_adver = AdversarialLoss()
self.swap_size = swap_size
self.swap = transforms.Randomswap((2, 2))
def __init__(self,
datamanager,
model,
optimizer,
margin=0.3,
weight_t=1,
weight_x=1,
scheduler=None,
use_gpu=True,
label_smooth=True):
super(ImageTripletEngine, self).__init__(datamanager, use_gpu)
self.model = model
self.optimizer = optimizer
self.scheduler = scheduler
self.register_model('model', model, optimizer, scheduler)
self.weight_t = weight_t
self.weight_x = weight_x
self.criterion_t = TripletLoss(margin=margin)
self.criterion_x = CrossEntropyLoss(
num_classes=self.datamanager.num_train_pids,
use_gpu=self.use_gpu,
label_smooth=label_smooth)
def __init__(
self,
datamanager,
model,
optimizer,
margin=0.3,
weight_t=1,
weight_x=1,
scheduler=None,
use_gpu=True,
label_smooth=True,
mmd_only=True,
):
super(ImageMmdEngine, self).__init__(datamanager, model, optimizer, scheduler, use_gpu, mmd_only)
self.optimizer.zero_grad()
self.mmd_only = mmd_only ###
self.weight_t = weight_t
self.weight_x = weight_x
self.criterion_t = TripletLoss(margin=margin)
self.criterion_x = CrossEntropyLoss(
num_classes=self.datamanager.num_train_pids,
use_gpu=self.use_gpu,
label_smooth=label_smooth
)
self.criterion_mmd = MaximumMeanDiscrepancy(
instances=self.datamanager.train_loader.sampler.num_instances,
batch_size=self.datamanager.train_loader.batch_size,
global_only=False,
distance_only=True,
all=False
)
def __init__(
self,
datamanager,
model,
optimizer,
margin=0.3,
weight_t=1,
weight_x=1,
weight_r=0.0001,
scheduler=None,
use_gpu=True,
label_smooth=True
):
super(ImageJointReconsVarEngine, self
).__init__(datamanager, model, optimizer, scheduler, use_gpu)
self.weight_t = weight_t
self.weight_x = weight_x
self.weight_r = weight_r
self.criterion_t = TripletLoss(margin=margin)
self.local_triplet = TripletLoss_Local(margin=margin)
self.done_once = True
self.criterion_x = CrossEntropyLoss(
num_classes=self.datamanager.num_train_pids,
use_gpu=self.use_gpu,
label_smooth=label_smooth
)
self.criterion_mse = torch.nn.MSELoss()
self.random = RandomErasing(probability=0.5,sl=0.02)
self.criterion = torch.nn.CrossEntropyLoss()
def __init__(self,
datamanager,
model,
optimizer,
margin=0.3,
weight_t=1,
weight_x=1,
weight_db_t=1,
weight_db_x=1,
weight_b_db_t=1,
weight_b_db_x=1,
scheduler=None,
use_gpu=True,
label_smooth=True,
top_drop_epoch=-1):
super(ImageTripletDropBatchDropBotFeaturesEngine,
self).__init__(datamanager, model, optimizer, scheduler, use_gpu)
self.weight_t = weight_t
self.weight_x = weight_x
self.weight_db_t = weight_db_t
self.weight_db_x = weight_db_x
self.weight_b_db_t = weight_b_db_t
self.weight_b_db_x = weight_b_db_x
self.top_drop_epoch = top_drop_epoch
self.criterion_t = TripletLoss(margin=margin)
self.criterion_x = CrossEntropyLoss(
num_classes=self.datamanager.num_train_pids,
use_gpu=self.use_gpu,
label_smooth=label_smooth)
self.criterion_db_t = TripletLoss(margin=margin)
self.criterion_db_x = CrossEntropyLoss(
num_classes=self.datamanager.num_train_pids,
use_gpu=self.use_gpu,
label_smooth=label_smooth)
self.criterion_b_db_t = TripletLoss(margin=margin)
self.criterion_b_db_x = CrossEntropyLoss(
num_classes=self.datamanager.num_train_pids,
use_gpu=self.use_gpu,
label_smooth=label_smooth)
def __init__(self, datamanager, model, optimizer, margin=0.3,
weight_t=1, weight_x=1, scheduler=None, use_cpu=False,
label_smooth=True, experiment=None, combine_method="mean", save_embed=None):
super(ImageTripletEngine, self).__init__(datamanager, model, optimizer, scheduler, use_cpu, experiment, combine_method, save_embed)
self.weight_t = weight_t
self.weight_x = weight_x
self.criterion_t = TripletLoss(margin=margin)
self.criterion_x = CrossEntropyLoss(
num_classes=self.datamanager.num_train_pids,
use_gpu=self.use_gpu,
label_smooth=label_smooth
)
def __init__(self, datamanager, model, optimizer, margin=0.3,
weight_t=1, weight_x=1, scheduler=None, use_gpu=True,
label_smooth=True):
super(ImageTripletEngine, self).__init__(datamanager, model, optimizer, scheduler, use_gpu)
self.weight_t = weight_t
self.weight_x = weight_x
self.criterion_t = TripletLoss(margin=margin)
self.criterion_x = CrossEntropyLoss(
num_classes=self.datamanager.num_train_pids,
use_gpu=self.use_gpu,
label_smooth=label_smooth
)
self.criterion_c = CenterLoss(num_classes=751, feat_dim=2048)
def __init__(
self,
datamanager,
model,
optimizer,
margin=0.27,
weight_t=1,
weight_x=1,
weight_r = 0.0000000001, #lambda
scheduler=None,
use_gpu=True,
label_smooth=True,
mmd_only=True,
datamanager2=None,
):
super(ImageMmdAEEngine, self).__init__(datamanager, model, optimizer, scheduler, use_gpu, mmd_only,datamanager2)
self.optimizer.zero_grad()
self.mmd_only = mmd_only ###
self.weight_t = weight_t
self.weight_x = weight_x
self.weight_r = weight_r
self.criterion_t = TripletLoss(margin=margin)
self.criterion_x = CrossEntropyLoss(
num_classes=self.datamanager.num_train_pids,
use_gpu=self.use_gpu,
label_smooth=label_smooth
)
self.criterion_mmd = MaximumMeanDiscrepancy(
instances=self.datamanager.train_loader.sampler.num_instances,
batch_size=self.datamanager.train_loader.batch_size,
global_only=False,
distance_only=True,
all=False,
)
self.criterion_mse = torch.nn.MSELoss()
self.random = RandomErasing(probability=0.5,sl=0.07)
self.randomt = RandomErasing(probability=0.5,sl=0.01)
self.mgn_loss = Loss(num_classes=self.datamanager.num_train_pids,use_gpu=self.use_gpu,label_smooth=label_smooth)
self.mgn_targetPredict =FC_Model().cuda()
self.BCE_criterion = torch.nn.BCEWithLogitsLoss()
def __init__(
self,
datamanager,
model1,
optimizer1,
scheduler1,
model2,
optimizer2,
scheduler2,
margin=0.3,
weight_t=0.5,
weight_x=1.,
weight_ml=1.,
use_gpu=True,
label_smooth=True,
deploy='model1'
):
super(ImageDMLEngine, self).__init__(datamanager, use_gpu)
self.model1 = model1
self.optimizer1 = optimizer1
self.scheduler1 = scheduler1
self.register_model('model1', model1, optimizer1, scheduler1)
self.model2 = model2
self.optimizer2 = optimizer2
self.scheduler2 = scheduler2
self.register_model('model2', model2, optimizer2, scheduler2)
self.weight_t = weight_t
self.weight_x = weight_x
self.weight_ml = weight_ml
assert deploy in ['model1', 'model2', 'both']
self.deploy = deploy
self.criterion_t = TripletLoss(margin=margin)
self.criterion_x = CrossEntropyLoss(
num_classes=self.datamanager.num_train_pids,
use_gpu=self.use_gpu,
label_smooth=label_smooth
)
def __init__(self,
datamanager,
model,
optimizer,
margin=0.3,
weight_t=1,
weight_x=1,
scheduler=None,
use_cpu=False,
label_smooth=True):
super(PoseTripleEngine, self).__init__(datamanager, model, optimizer,
scheduler, use_cpu)
self.weight_t = weight_t
self.weight_x = weight_x
self.criterion_t = TripletLoss(margin=margin)
# TODO modify the criterion for pairwise comparison
self.criterion_x = CrossEntropyLoss(
num_classes=self.datamanager.num_train_pids,
use_gpu=self.use_gpu,
label_smooth=label_smooth)
self.att_criterion = Isolate_loss()
def __init__(
self,
datamanager,
model,
optimizer,
margin=0.3,
weight_t=1,
weight_x=1,
weight_r=0.0001,
scheduler=None,
use_gpu=True,
label_smooth=True,
only_recons=True,
):
super(ImageReconsVarEngine,
self).__init__(datamanager, model, optimizer, scheduler, use_gpu)
self.optimizer.zero_grad()
self.weight_t = weight_t
self.weight_x = weight_x
self.weight_r = weight_r
self.criterion_t = TripletLoss(margin=margin)
self.criterion_x = CrossEntropyLoss(
num_classes=self.datamanager.num_train_pids,
use_gpu=self.use_gpu,
label_smooth=label_smooth)
self.criterion_mse = torch.nn.MSELoss()
self.criterion_mmd = MaximumMeanDiscrepancy(
instances=self.datamanager.train_loader.sampler.num_instances,
batch_size=self.datamanager.train_loader.batch_size,
global_only=False,
distance_only=True,
all=False)
self.only_recons = only_recons
self.random = RandomErasing(probability=0.5)
self.random2 = RandomErasing(probability=0.65, sl=0.15)
def __init__(self,
datamanager,
model,
optimizer,
margin=0.3,
weight_t=1,
weight_x=1,
weight_r=1,
scheduler=None,
use_gpu=True,
label_smooth=True):
super(ImageJointReconsEngine,
self).__init__(datamanager, model, optimizer, scheduler, use_gpu)
self.weight_t = weight_t
self.weight_x = weight_x
self.weight_r = weight_r
self.criterion_t = TripletLoss(margin=margin)
self.criterion_x = CrossEntropyLoss(
num_classes=self.datamanager.num_train_pids,
use_gpu=self.use_gpu,
label_smooth=label_smooth)
self.criterion_mse = torch.nn.MSELoss()
def __init__(self, num_classes,use_gpu,label_smooth):
super(Loss, self).__init__()
self.cross_entropy_loss = CrossEntropyLossTorch()
self.triplet_loss = TripletLoss(margin=1.2)
def main():
global best_rank1, best_mAP
random.seed(args.seed)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
if not args.use_avai_gpus:
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices
use_gpu = torch.cuda.is_available()
if args.use_cpu:
use_gpu = False
if not args.evaluate:
sys.stdout = Logger(
osp.join(
args.save_dir,
'log_train{}.txt'.format(time.strftime('-%Y-%m-%d-%H-%M-%S'))))
else:
sys.stdout = Logger(
osp.join(
args.save_dir,
'log_test{}.txt'.format(time.strftime('-%Y-%m-%d-%H-%M-%S'))))
writer = SummaryWriter(log_dir=args.save_dir, comment=args.arch)
print("==========\nArgs:{}\n==========".format(args))
if use_gpu:
print("Currently using GPU {}".format(args.gpu_devices))
cudnn.benchmark = False if 'resnet3dt' in args.arch else True
torch.cuda.manual_seed_all(args.seed)
else:
print("Currently using CPU (GPU is highly recommended)")
print("Initializing dataset {}".format(args.dataset))
dataset = data_manager.init_vidreid_dataset(root=args.root,
name=args.dataset,
split_id=args.split_id,
use_pose=args.use_pose)
transform_train = list()
print('Transform:')
if args.misalign_aug:
print('+ Misalign Augmentation')
transform_train.append(T.GroupMisAlignAugment())
if args.rand_crop:
print('+ Random Crop')
transform_train.append(T.GroupRandomCrop(size=(240, 120)))
print('+ Resize to ({} x {})'.format(args.height, args.width))
transform_train.append(T.GroupResize((args.height, args.width)))
if args.flip_aug:
print('+ Random HorizontalFlip')
transform_train.append(T.GroupRandomHorizontalFlip())
print('+ ToTensor')
transform_train.append(T.GroupToTensor())
print(
'+ Normalize with mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]'
)
transform_train.append(
T.GroupNormalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224,
0.225]))
if args.rand_erase:
print('+ Random Erasing')
transform_train.append(T.GroupRandomErasing())
transform_train = T.Compose(transform_train)
transform_test = T.Compose([
T.GroupResize((args.height, args.width)),
T.GroupToTensor(),
T.GroupNormalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224,
0.225]),
])
pin_memory = True if use_gpu else False
trainloader = DataLoader(
VideoDataset(dataset.train,
seq_len=args.seq_len,
sample=args.train_sample,
transform=transform_train,
training=True,
pose_info=dataset.process_poses,
num_split=args.num_split,
num_parts=args.num_parts,
num_scale=args.num_scale,
pyramid_part=args.pyramid_part,
enable_pose=args.use_pose),
sampler=eval(args.train_sampler)(dataset.train,
batch_size=args.train_batch,
num_instances=args.num_instances),
batch_size=args.train_batch,
num_workers=args.workers,
pin_memory=pin_memory,
drop_last=True,
)
queryloader = DataLoader(
VideoDataset(dataset.query,
seq_len=args.seq_len,
sample=args.test_sample,
transform=transform_test,
pose_info=dataset.process_poses,
num_split=args.num_split,
num_parts=args.num_parts,
num_scale=args.num_scale,
pyramid_part=args.pyramid_part,
enable_pose=args.use_pose),
batch_size=args.test_batch,
shuffle=False,
num_workers=args.workers,
pin_memory=pin_memory,
drop_last=False,
)
galleryloader = DataLoader(
VideoDataset(dataset.gallery,
seq_len=args.seq_len,
sample=args.test_sample,
transform=transform_test,
pose_info=dataset.process_poses,
num_split=args.num_split,
num_parts=args.num_parts,
num_scale=args.num_scale,
pyramid_part=args.pyramid_part,
enable_pose=args.use_pose),
batch_size=args.test_batch,
shuffle=False,
num_workers=args.workers,
pin_memory=pin_memory,
drop_last=False,
)
print("Initializing model: {}".format(args.arch))
model = models.init_model(name=args.arch,
num_classes=dataset.num_train_pids,
loss={'xent', 'htri'},
last_stride=args.last_stride,
num_parts=args.num_parts,
num_scale=args.num_scale,
num_split=args.num_split,
pyramid_part=args.pyramid_part,
num_gb=args.num_gb,
use_pose=args.use_pose,
learn_graph=args.learn_graph,
consistent_loss=args.consistent_loss,
bnneck=args.bnneck,
save_dir=args.save_dir)
input_size = sum(calc_splits(
args.num_split)) if args.pyramid_part else args.num_split
input_size *= args.num_scale * args.seq_len
num_params, flops = compute_model_complexity(
model,
input=[
torch.randn(1, args.seq_len, 3, args.height, args.width),
torch.ones(1, input_size, input_size)
],
verbose=True,
only_conv_linear=False)
print('Model complexity: params={:,} flops={:,}'.format(num_params, flops))
if args.label_smooth:
criterion_xent = CrossEntropyLabelSmooth(
num_classes=dataset.num_train_pids, use_gpu=use_gpu)
else:
criterion_xent = nn.CrossEntropyLoss()
criterion_htri = TripletLoss(margin=args.margin, soft=args.soft_margin)
param_groups = model.parameters()
optimizer = init_optim(args.optim, param_groups, args.lr,
args.weight_decay)
scheduler = lr_scheduler.MultiStepLR(optimizer,
milestones=args.stepsize,
gamma=args.gamma)
if args.warmup:
scheduler = lr_scheduler.WarmupMultiStepLR(optimizer,
milestones=args.stepsize,
gamma=args.gamma,
warmup_iters=10,
warmup_factor=0.01)
if args.load_weights and check_isfile(args.load_weights):
# load pretrained weights but ignore layers that don't match in size
checkpoint = torch.load(args.load_weights)
pretrain_dict = checkpoint['state_dict']
model_dict = model.state_dict()
pretrain_dict = {
k: v
for k, v in pretrain_dict.items()
if k in model_dict and model_dict[k].size() == v.size()
}
model_dict.update(pretrain_dict)
model.load_state_dict(model_dict)
print("Loaded pretrained weights from '{}'".format(args.load_weights))
if args.resume and check_isfile(args.resume):
print("Loaded checkpoint from '{}'".format(args.resume))
from functools import partial
import pickle
pickle.load = partial(pickle.load, encoding="latin1")
pickle.Unpickler = partial(pickle.Unpickler, encoding="latin1")
checkpoint = torch.load(args.resume,
map_location=lambda storage, loc: storage,
pickle_module=pickle)
print('Loaded model weights')
model.load_state_dict(checkpoint['state_dict'])
if optimizer is not None and 'optimizer' in checkpoint:
print('Loaded optimizer')
optimizer.load_state_dict(checkpoint['optimizer'])
if use_gpu:
for state in optimizer.state.values():
for k, v in state.items():
if isinstance(v, torch.Tensor):
state[k] = v.cuda()
start_epoch = checkpoint['epoch'] + 1
print('- start_epoch: {}'.format(start_epoch))
best_rank1 = checkpoint['rank1']
print("- rank1: {}".format(best_rank1))
if 'mAP' in checkpoint:
best_mAP = checkpoint['mAP']
print("- mAP: {}".format(best_mAP))
else:
start_epoch = 0
if use_gpu:
model = nn.DataParallel(model).cuda()
if args.evaluate:
print("Evaluate only")
distmat = test(model,
queryloader,
galleryloader,
args.pool,
use_gpu,
return_distmat=True)
if args.visualize_ranks:
visualize_ranked_results(
distmat,
dataset,
save_dir=osp.join(args.save_dir, 'ranked_results'),
topk=20,
)
return
start_time = time.time()
train_time = 0
best_epoch = start_epoch
print("==> Start training")
for epoch in range(start_epoch, args.max_epoch):
start_train_time = time.time()
train(epoch,
model,
criterion_xent,
criterion_htri,
optimizer,
trainloader,
use_gpu,
writer=writer)
train_time += round(time.time() - start_train_time)
if epoch >= args.zero_wd > 0:
set_wd(optimizer, 0)
for group in optimizer.param_groups:
assert group['weight_decay'] == 0, '{} is not zero'.format(
group['weight_decay'])
scheduler.step(epoch)
if (epoch + 1) > args.start_eval and args.eval_step > 0 and (
epoch + 1) % args.eval_step == 0 or (epoch +
1) == args.max_epoch:
print("==> Test")
rank1, mAP = test(model, queryloader, galleryloader, args.pool,
use_gpu)
is_best = rank1 > best_rank1
if is_best:
best_rank1 = rank1
best_mAP = mAP
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,
'optimizer': optimizer.state_dict(),
'rank1': rank1,
'mAP': mAP,
'epoch': epoch,
}, False,
osp.join(args.save_dir,
'checkpoint_ep' + str(epoch + 1) + '.pth.tar'))
writer.add_scalar(tag='acc/rank1',
scalar_value=rank1,
global_step=epoch + 1)
writer.add_scalar(tag='acc/mAP',
scalar_value=mAP,
global_step=epoch + 1)
print("==> Best Rank-1 {:.2%}, mAP: {:.2%}, achieved at epoch {}".format(
best_rank1, best_mAP, best_epoch))
elapsed = round(time.time() - start_time)
elapsed = str(datetime.timedelta(seconds=elapsed))
train_time = str(datetime.timedelta(seconds=train_time))
print(
"Finished. Total elapsed time (h:m:s): {}. Training time (h:m:s): {}.".
format(elapsed, train_time))
def main():
global args
torch.manual_seed(args.seed)
if not args.use_avai_gpus: os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices
use_gpu = torch.cuda.is_available()
if args.use_cpu: use_gpu = False
log_name = 'log_test.txt' if args.evaluate else 'log_train.txt'
sys.stdout = Logger(osp.join(args.save_dir, log_name))
print("==========\nArgs:{}\n==========".format(args))
if use_gpu:
print("Currently using GPU {}".format(args.gpu_devices))
cudnn.benchmark = True
torch.cuda.manual_seed_all(args.seed)
else:
print("Currently using CPU, however, GPU is highly recommended")
print("Initializing image data manager")
dm = ImageDataManager(use_gpu, **image_dataset_kwargs(args))
trainloader, testloader_dict = dm.return_dataloaders()
print("Initializing model: {}".format(args.arch))
model = models.init_model(name=args.arch, num_classes=dm.num_train_pids, loss={'xent', 'htri'})
print("Model size: {:.3f} M".format(count_num_param(model)))
criterion_xent = CrossEntropyLoss(num_classes=dm.num_train_pids, use_gpu=use_gpu, label_smooth=args.label_smooth)
criterion_htri = TripletLoss(margin=args.margin)
optimizer = init_optimizer(model.parameters(), **optimizer_kwargs(args))
scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=args.stepsize, gamma=args.gamma)
if args.load_weights and check_isfile(args.load_weights):
# load pretrained weights but ignore layers that don't match in size
checkpoint = torch.load(args.load_weights)
pretrain_dict = checkpoint['state_dict']
model_dict = model.state_dict()
pretrain_dict = {k: v for k, v in pretrain_dict.items() if k in model_dict and model_dict[k].size() == v.size()}
model_dict.update(pretrain_dict)
model.load_state_dict(model_dict)
print("Loaded pretrained weights from '{}'".format(args.load_weights))
if args.resume and check_isfile(args.resume):
checkpoint = torch.load(args.resume)
model.load_state_dict(checkpoint['state_dict'])
args.start_epoch = checkpoint['epoch'] + 1
print("Loaded checkpoint from '{}'".format(args.resume))
print("- start_epoch: {}\n- rank1: {}".format(args.start_epoch, checkpoint['rank1']))
if use_gpu:
model = nn.DataParallel(model).cuda()
if args.evaluate:
print("Evaluate only")
for name in args.target_names:
print("Evaluating {} ...".format(name))
queryloader = testloader_dict[name]['query']
galleryloader = testloader_dict[name]['gallery']
distmat = test(model, queryloader, galleryloader, use_gpu, return_distmat=True)
if args.visualize_ranks:
visualize_ranked_results(
distmat, dm.return_testdataset_by_name(name),
save_dir=osp.join(args.save_dir, 'ranked_results', name),
topk=20
)
return
start_time = time.time()
ranklogger = RankLogger(args.source_names, args.target_names)
train_time = 0
print("=> Start training")
if args.fixbase_epoch > 0:
print("Train {} for {} epochs while keeping other layers frozen".format(args.open_layers, args.fixbase_epoch))
initial_optim_state = optimizer.state_dict()
for epoch in range(args.fixbase_epoch):
start_train_time = time.time()
train(epoch, model, criterion_xent, criterion_htri, optimizer, trainloader, use_gpu, fixbase=True)
train_time += round(time.time() - start_train_time)
print("Done. All layers are open to train for {} epochs".format(args.max_epoch))
optimizer.load_state_dict(initial_optim_state)
for epoch in range(args.start_epoch, args.max_epoch):
start_train_time = time.time()
train(epoch, model, criterion_xent, criterion_htri, optimizer, trainloader, use_gpu)
train_time += round(time.time() - start_train_time)
scheduler.step()
if (epoch + 1) > args.start_eval and args.eval_freq > 0 and (epoch + 1) % args.eval_freq == 0 or (epoch + 1) == args.max_epoch:
print("=> Test")
for name in args.target_names:
print("Evaluating {} ...".format(name))
queryloader = testloader_dict[name]['query']
galleryloader = testloader_dict[name]['gallery']
rank1 = test(model, queryloader, galleryloader, use_gpu)
ranklogger.write(name, epoch + 1, rank1)
if use_gpu:
state_dict = model.module.state_dict()
else:
state_dict = model.state_dict()
save_checkpoint({
'state_dict': state_dict,
'rank1': rank1,
'epoch': epoch,
}, False, osp.join(args.save_dir, 'checkpoint_ep' + str(epoch + 1) + '.pth.tar'))
elapsed = round(time.time() - start_time)
elapsed = str(datetime.timedelta(seconds=elapsed))
train_time = str(datetime.timedelta(seconds=train_time))
print("Finished. Total elapsed time (h:m:s): {}. Training time (h:m:s): {}.".format(elapsed, train_time))
ranklogger.show_summary()
def main():
global args
torch.manual_seed(args.seed)
if not args.use_avai_gpus:
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices
use_gpu = torch.cuda.is_available()
if args.use_cpu: use_gpu = False
log_name = 'log_test.txt' if args.evaluate else 'log_train.txt'
sys.stdout = Logger(osp.join(args.save_dir, log_name))
print("==========\nArgs:{}\n==========".format(args))
if use_gpu:
print("Currently using GPU {}".format(args.gpu_devices))
cudnn.benchmark = True
torch.cuda.manual_seed_all(args.seed)
else:
print("Currently using CPU, however, GPU is highly recommended")
print("Initializing image data manager")
dm = ImageDataManager(use_gpu, **image_dataset_kwargs(args))
trainloader, testloader_dict = dm.return_dataloaders()
# ReID-Stream:
print("Initializing ReID-Stream: {}".format(args.arch))
model = models.init_model(name=args.arch,
num_classes=dm.num_train_pids,
reid_dim=args.reid_dim,
loss={'xent', 'htri'})
print("ReID Model size: {:.3f} M".format(count_num_param(model)))
criterion_xent = CrossEntropyLoss(num_classes=dm.num_train_pids,
use_gpu=use_gpu,
label_smooth=args.label_smooth)
criterion_htri = TripletLoss(margin=args.margin)
# 2. Optimizer
# Main ReID-Stream:
optimizer = init_optimizer(model.parameters(), **optimizer_kwargs(args))
scheduler = lr_scheduler.MultiStepLR(optimizer,
milestones=args.stepsize,
gamma=args.gamma)
if use_gpu:
model = nn.DataParallel(model).cuda()
if args.evaluate:
print("Evaluate only")
for name in args.target_names:
print("Evaluating {} ...".format(name))
queryloader = testloader_dict[name]['query']
galleryloader = testloader_dict[name]['gallery']
distmat = test(model,
queryloader,
galleryloader,
use_gpu,
return_distmat=True)
if args.visualize_ranks:
visualize_ranked_results(distmat,
dm.return_testdataset_by_name(name),
save_dir=osp.join(
args.save_dir, 'ranked_results',
name),
topk=20)
return
start_time = time.time()
ranklogger = RankLogger(args.source_names, args.target_names)
train_time = 0
print("==> Start training")
for epoch in range(args.start_epoch, args.max_epoch):
start_train_time = time.time()
train(epoch, model, criterion_xent, criterion_htri, \
optimizer, trainloader, use_gpu)
train_time += round(time.time() - start_train_time)
scheduler.step()
if (epoch + 1) > args.start_eval and args.eval_freq > 0 and (
epoch + 1) % args.eval_freq == 0 or (epoch +
1) == args.max_epoch:
print("==> Test")
for name in args.target_names:
print("Evaluating {} ...".format(name))
queryloader = testloader_dict[name]['query']
galleryloader = testloader_dict[name]['gallery']
rank1 = test(model, queryloader, galleryloader, use_gpu)
ranklogger.write(name, epoch + 1, rank1)
if use_gpu:
state_dict = model.module.state_dict()
else:
state_dict = model.state_dict()
save_checkpoint(
{
'state_dict': state_dict,
'rank1': rank1,
'epoch': epoch,
}, False,
osp.join(args.save_dir,
'checkpoint_ep' + str(epoch + 1) + '.pth.tar'))
elapsed = round(time.time() - start_time)
elapsed = str(datetime.timedelta(seconds=elapsed))
train_time = str(datetime.timedelta(seconds=train_time))
print(
"Finished. Total elapsed time (h:m:s): {}. Training time (h:m:s): {}.".
format(elapsed, train_time))
ranklogger.show_summary()
def main():
parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('--config-file', type=str, default='', help='path to config file')
parser.add_argument('--gpu-devices', type=str, default='', )
parser.add_argument('opts', default=None, nargs=argparse.REMAINDER,
help='Modify config options using the command-line')
args = parser.parse_args()
cfg = get_default_config()
if args.config_file:
cfg.merge_from_file(args.config_file)
cfg.merge_from_list(args.opts)
set_random_seed(cfg.train.seed)
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices
log_name = 'test.log' if cfg.test.evaluate else 'train.log'
log_name += time.strftime('-%Y-%m-%d-%H-%M-%S')
sys.stdout = Logger(osp.join(cfg.data.save_dir, log_name))
print('Show configuration\n{}\n'.format(cfg))
torch.backends.cudnn.benchmark = True
datamanager = ImageDataManager(**imagedata_kwargs(cfg))
trainloader, queryloader, galleryloader = datamanager.return_dataloaders()
print('Building model: {}'.format(cfg.model.name))
model = build_model(cfg.model.name, datamanager.num_train_pids, 'triplet', pretrained=cfg.model.pretrained)
if cfg.model.load_weights and check_isfile(cfg.model.load_weights):
load_pretrained_weights(model, cfg.model.load_weights)
model = nn.DataParallel(model).cuda()
criterion_t = TripletLoss(margin=cfg.loss.triplet.margin)
criterion_x = CrossEntropyLoss(datamanager.num_train_pids, label_smooth=cfg.loss.softmax.label_smooth)
optimizer = build_optimizer(model, **optimizer_kwargs(cfg))
scheduler = build_lr_scheduler(optimizer, **lr_scheduler_kwargs(cfg))
if cfg.model.resume and check_isfile(cfg.model.resume):
cfg.train.start_epoch = resume_from_checkpoint(cfg.model.resume, model, optimizer=optimizer)
if cfg.test.evaluate:
distmat = evaluate(model, queryloader, galleryloader, dist_metric=cfg.test.dist_metric,
normalize_feature=cfg.test.normalize_feature, rerank=cfg.test.rerank, return_distmat=True)
if cfg.test.visrank:
visualize_ranked_results(distmat, datamanager.return_testdataset(), 'image', width=cfg.data.width,
height=cfg.data.height, save_dir=osp.join(cfg.data.save_dir, 'visrank'))
return
time_start = time.time()
print('=> Start training')
for epoch in range(cfg.train.start_epoch, cfg.train.max_epoch):
train(epoch, cfg.train.max_epoch, model, criterion_t, criterion_x, optimizer, trainloader,
fixbase_epoch=cfg.train.fixbase_epoch, open_layers=cfg.train.open_layers)
scheduler.step()
if (epoch + 1) % cfg.test.eval_freq == 0 or (epoch + 1) == cfg.train.max_epoch:
rank1 = evaluate(model, queryloader, galleryloader, dist_metric=cfg.test.dist_metric,
normalize_feature=cfg.test.normalize_feature, rerank=cfg.test.rerank)
save_checkpoint({
'state_dict': model.state_dict(),
'epoch': epoch + 1,
'rank1': rank1,
'optimizer': optimizer.state_dict(),
}, cfg.data.save_dir)
elapsed = round(time.time() - time_start)
elapsed = str(datetime.timedelta(seconds=elapsed))
print('Elapsed {}'.format(elapsed))
def main():
################################# Setting gpu and out file ###################################################################
if args.use_cpu: use_gpu = False
use_gpu = 1
if use_gpu == 1:
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices
os.environ['PYTHONHASHSEED'] = str(args.seed)
if args.is_deterministic == 1:
print("Currently using GPU {}".format(args.gpu_devices))
print('is deterministic')
np.random.seed(args.seed)
torch.manual_seed(args.seed)
random.seed(args.seed)
torch.cuda.manual_seed(args.seed)
torch.cuda.manual_seed_all(args.seed)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
use_gpu = torch.cuda.is_available()
else:
print("Currently using CPU (GPU is highly recommended)")
if not args.evaluate:
sys.stdout = Logger(osp.join(args.save_dir, 'log_train.txt'))
else:
sys.stdout = Logger(osp.join(args.save_dir, 'log_test.txt'))
print("==========\nArgs:{}\n==========".format(args))
print("Initializing dataset {}".format(args.dataset))
[dataset, trainloader_reid, queryloader,
galleryloader] = initialize_loader(use_gpu)
################################### Initialize model ###############################################################
if args.num_classes_attributes == 0:
num_classes_attributes = (6, 4, 11, 14, 6)
if args.num_classes_attributes > 1:
num_classes_attributes = (33, 6, 58, 65, 9)
if args.dataset == 'pa100K':
num_classes_attributes = (28, 9, 3, 14, 24, 9)
if args.num_classes_attributes == 0:
num_classes_attributes = (5, 3, 3, 6, 9, 10)
if args.dataset == 'dukemtmcreid':
num_classes_attributes = (15, 2, 47, 62, 9)
if args.num_classes_attributes == 0:
num_classes_attributes = (2, 2, 8, 11, 6)
num_group_attributes = len(num_classes_attributes)
print("Initializing model: {}".format(args.arch))
if args.arch == 'resnetAttW2VText':
dim_features = 50 * len(num_classes_attributes)
model = models.init_model(args.arch,
num_classes=dataset.num_train_pids,
loss={'xent', 'htri'},
num_group_attributes=num_group_attributes,
dim_features=50,
glove=True)
criterion_attributes = LossW2V(
num_classes_attributes=num_classes_attributes,
attr_loss_type=args.attr_loss_type)
if args.arch == 'resnetAttW2VAttributes':
dim_features = sum(num_classes_attributes)
model = models.init_model(
args.arch,
num_classes=dataset.num_train_pids,
loss={'xent', 'htri'},
num_group_attributes=num_group_attributes,
num_classes_attributes=num_classes_attributes,
dim_features=dim_features,
glove=True)
criterion_attributes = LossW2V(
num_classes_attributes=num_classes_attributes)
all_parameters = model.parameters()
optimizer = init_optim(args.optim, all_parameters, args.lr,
args.weight_decay)
################################### Loss functions ##############################
print("Model size: {:.3f} M".format(count_num_param(model)))
criterion_htri_reid = TripletLoss(margin=0.3)
criterion_xent_reid = nn.CrossEntropyLoss()
################################### Pretrained models ##############################
if args.load_weights:
load_weights(model)
if args.resume:
resume(model)
if use_gpu:
model = nn.DataParallel(model).cuda()
################################### Only evaluation ###############################
writer = SummaryWriter(log_dir=osp.join(args.exp_dir, 'tensorboard'))
if args.evaluate:
print("Evaluate by test")
test_and_evaluate_dist_mat(writer,
model,
queryloader,
galleryloader,
use_gpu=use_gpu,
save_features=True,
load_features=False,
arch=args.arch,
size=dim_features)
return
################################### Training ######################################
start_time = time.time()
train_time = 0
best_rank1 = -np.inf
best_epoch = 0
# Schedule learning rate
print("==> Start training")
for epoch in range(args.start_epoch, args.max_epoch):
start_train_time = time.time()
if args.is_warmup:
adjust_lr(epoch, optimizer, args.lr)
if args.arch == 'resnetAttW2VText':
train_w2v_single_batch_text(writer, epoch, model,
criterion_htri_reid,
criterion_attributes, optimizer,
trainloader_reid, use_gpu)
if args.arch == 'resnetAttW2VAttributes':
train_w2v_single_batch_attributes(writer, epoch, model,
criterion_htri_reid,
criterion_attributes, optimizer,
trainloader_reid, use_gpu)
train_time += round(time.time() - start_train_time)
if (epoch + 1) > args.start_eval and args.eval_step > 0 and (
epoch + 1) % args.eval_step == 0 or (epoch +
1) == args.max_epoch:
print("==> Test")
elapsed = round(time.time() - start_time)
elapsed = str(datetime.timedelta(seconds=elapsed))
train_time = str(datetime.timedelta(seconds=train_time))
print(
"Finished. Total elapsed time (h:m:s): {}. Training time (h:m:s): {}.".
format(elapsed, train_time))
state_dict = model.module.state_dict()
save_checkpoint(state={
'state_dict': state_dict,
'epoch': epoch,
},
fpath=osp.join(
args.save_dir,
'checkpoint_ep' + str(epoch + 1) + '.pth.tar'))
print("Evaluate by test")
test_and_evaluate_dist_mat(writer,
model,
queryloader,
galleryloader,
use_gpu=use_gpu,
save_features=True,
load_features=False,
arch=args.arch,
size=dim_features)
def main():
torch.manual_seed(args.seed)
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices
use_gpu = torch.cuda.is_available()
if args.use_cpu: use_gpu = False
if not args.evaluate:
sys.stdout = Logger(osp.join(args.save_dir, 'log_train.txt'))
else:
sys.stdout = Logger(osp.join(args.save_dir, 'log_test.txt'))
print("==========\nArgs:{}\n==========".format(args))
if use_gpu:
print("Currently using GPU {}".format(args.gpu_devices))
cudnn.benchmark = True
torch.cuda.manual_seed_all(args.seed)
else:
print("Currently using CPU (GPU is highly recommended)")
print("Initializing dataset {}".format(args.dataset))
dataset = data_manager.init_vidreid_dataset(root=args.root, name=args.dataset)
transform_train = T.Compose([
T.Random2DTranslation(args.height, args.width),
T.RandomHorizontalFlip(),
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
transform_test = T.Compose([
T.Resize((args.height, args.width)),
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
pin_memory = True if use_gpu else False
# decompose tracklets into images for image-based training
new_train = []
for img_paths, pid, camid in dataset.train:
for img_path in img_paths:
new_train.append((img_path, pid, camid))
trainloader = DataLoader(
ImageDataset(new_train, transform=transform_train),
sampler=RandomIdentitySampler(new_train, args.train_batch, args.num_instances),
batch_size=args.train_batch, num_workers=args.workers,
pin_memory=pin_memory, drop_last=True,
)
queryloader = DataLoader(
VideoDataset(dataset.query, seq_len=args.seq_len, sample='evenly', transform=transform_test),
batch_size=args.test_batch, shuffle=False, num_workers=args.workers,
pin_memory=pin_memory, drop_last=False,
)
galleryloader = DataLoader(
VideoDataset(dataset.gallery, seq_len=args.seq_len, sample='evenly', transform=transform_test),
batch_size=args.test_batch, shuffle=False, num_workers=args.workers,
pin_memory=pin_memory, drop_last=False,
)
print("Initializing model: {}".format(args.arch))
model = models.init_model(name=args.arch, num_classes=dataset.num_train_pids, loss={'xent', 'htri'})
print("Model size: {:.3f} M".format(count_num_param(model)))
criterion_xent = nn.CrossEntropyLoss()
criterion_htri = TripletLoss(margin=args.margin)
criterion_KA = KALoss(margin=args.margin, same_margin = args.same_margin, use_auto_samemargin = args.use_auto_samemargin)
cirterion_lifted = LiftedLoss(margin=args.margin)
cirterion_batri = BA_TripletLoss(margin=args.margin)
if args.use_auto_samemargin == True:
G_params = [{'params': model.parameters(), 'lr': args.lr }, {'params': criterion_KA.auto_samemargin, 'lr': args.lr}]
else :
G_params = [para for _, para in model.named_parameters()]
optimizer = init_optim(args.optim, G_params, args.lr, args.weight_decay)
if args.load_weights:
# load pretrained weights but ignore layers that don't match in size
if check_isfile(args.load_weights):
checkpoint = torch.load(args.load_weights)
pretrain_dict = checkpoint['state_dict']
model_dict = model.state_dict()
pretrain_dict = {k: v for k, v in pretrain_dict.items() if k in model_dict and model_dict[k].size() == v.size()}
model_dict.update(pretrain_dict)
model.load_state_dict(model_dict)
print("Loaded pretrained weights from '{}'".format(args.load_weights))
if args.resume:
if check_isfile(args.resume):
checkpoint = torch.load(args.resume)
model.load_state_dict(checkpoint['state_dict'])
args.start_epoch = checkpoint['epoch']
rank1 = checkpoint['rank1']
print("Loaded checkpoint from '{}'".format(args.resume))
print("- start_epoch: {}\n- rank1: {}".format(args.start_epoch, rank1))
if use_gpu:
model = nn.DataParallel(model).cuda()
if args.evaluate:
print("Evaluate only")
distmat = test(model, queryloader, galleryloader, args.pool, use_gpu, return_distmat=True)
if args.vis_ranked_res:
visualize_ranked_results(
distmat, dataset,
save_dir=osp.join(args.save_dir, 'ranked_results'),
topk=20,
)
return
start_time = time.time()
train_time = 0
best_rank1 = -np.inf
best_epoch = 0
print("==> Start training")
for epoch in range(args.start_epoch, args.max_epoch):
start_train_time = time.time()
adjust_learning_rate(optimizer, epoch)
train(epoch, model, cirterion_batri, cirterion_lifted, criterion_xent, criterion_htri, criterion_KA, optimizer, trainloader, use_gpu)
train_time += round(time.time() - start_train_time)
scheduler.step()
if (epoch + 1) > args.start_eval and args.eval_step > 0 and (epoch + 1) % args.eval_step == 0 or (epoch + 1) == args.max_epoch:
print("==> Test")
rank1 = test(model, queryloader, galleryloader, args.pool, use_gpu)
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,
'rank1': rank1,
'epoch': epoch,
}, is_best, osp.join(args.save_dir, 'checkpoint_ep' + str(epoch + 1) + '.pth.tar'))
print("==> Best Rank-1 {:.1%}, achieved at epoch {}".format(best_rank1, best_epoch))
elapsed = round(time.time() - start_time)
elapsed = str(datetime.timedelta(seconds=elapsed))
train_time = str(datetime.timedelta(seconds=train_time))
print("Finished. Total elapsed time (h:m:s): {}. Training time (h:m:s): {}.".format(elapsed, train_time))
def main():
global args, best_rank1
torch.manual_seed(args.seed)
if not args.use_avai_gpus:
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices
use_gpu = torch.cuda.is_available()
if args.use_cpu: use_gpu = False
if not args.evaluate:
sys.stdout = Logger(osp.join(args.save_dir, 'log_train.txt'))
else:
sys.stdout = Logger(osp.join(args.save_dir, 'log_test.txt'))
print("==========\nArgs:{}\n==========".format(args))
if use_gpu:
print("Currently using GPU {}".format(args.gpu_devices))
cudnn.benchmark = True
torch.cuda.manual_seed_all(args.seed)
else:
print("Currently using CPU (GPU is highly recommended)")
print("Initializing dataset {}".format(args.dataset))
dataset = data_manager.init_imgreid_dataset(
root=args.root,
name=args.dataset,
split_id=args.split_id,
cuhk03_labeled=args.cuhk03_labeled,
cuhk03_classic_split=args.cuhk03_classic_split,
)
transform_train = T.Compose([
T.Random2DTranslation(args.height, args.width),
T.RandomHorizontalFlip(),
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
transform_test = T.Compose([
T.Resize((args.height, args.width)),
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
pin_memory = True if use_gpu else False
trainloader = DataLoader(
ImageDataset(dataset.train, transform=transform_train),
sampler=RandomIdentitySampler(dataset.train, args.train_batch,
args.num_instances),
batch_size=args.train_batch,
num_workers=args.workers,
pin_memory=pin_memory,
drop_last=True,
)
queryloader = DataLoader(
ImageDataset(dataset.query, transform=transform_test),
batch_size=args.test_batch,
shuffle=False,
num_workers=args.workers,
pin_memory=pin_memory,
drop_last=False,
)
galleryloader = DataLoader(
ImageDataset(dataset.gallery, transform=transform_test),
batch_size=args.test_batch,
shuffle=False,
num_workers=args.workers,
pin_memory=pin_memory,
drop_last=False,
)
print("Initializing model: {}".format(args.arch))
model = models.init_model(name=args.arch,
num_classes=dataset.num_train_pids,
loss={'xent', 'htri'})
print("Model size: {:.3f} M".format(count_num_param(model)))
if args.label_smooth:
criterion_xent = CrossEntropyLabelSmooth(
num_classes=dataset.num_train_pids, use_gpu=use_gpu)
else:
criterion_xent = nn.CrossEntropyLoss()
criterion_htri = TripletLoss(margin=args.margin)
optimizer = init_optim(args.optim, model.parameters(), args.lr,
args.weight_decay)
scheduler = lr_scheduler.MultiStepLR(optimizer,
milestones=args.stepsize,
gamma=args.gamma)
if args.load_weights and check_isfile(args.load_weights):
# load pretrained weights but ignore layers that don't match in size
checkpoint = torch.load(args.load_weights)
pretrain_dict = checkpoint['state_dict']
model_dict = model.state_dict()
pretrain_dict = {
k: v
for k, v in pretrain_dict.items()
if k in model_dict and model_dict[k].size() == v.size()
}
model_dict.update(pretrain_dict)
model.load_state_dict(model_dict)
print("Loaded pretrained weights from '{}'".format(args.load_weights))
if args.resume and check_isfile(args.resume):
checkpoint = torch.load(args.resume)
model.load_state_dict(checkpoint['state_dict'])
args.start_epoch = checkpoint['epoch'] + 1
best_rank1 = checkpoint['rank1']
print("Loaded checkpoint from '{}'".format(args.resume))
print("- start_epoch: {}\n- rank1: {}".format(args.start_epoch,
best_rank1))
if use_gpu:
model = nn.DataParallel(model).cuda()
if args.evaluate:
print("Evaluate only")
distmat = test(model,
queryloader,
galleryloader,
use_gpu,
return_distmat=True)
if args.visualize_ranks:
visualize_ranked_results(
distmat,
dataset,
save_dir=osp.join(args.save_dir, 'ranked_results'),
topk=20,
)
return
start_time = time.time()
train_time = 0
best_epoch = args.start_epoch
print("==> Start training")
for epoch in range(args.start_epoch, args.max_epoch):
start_train_time = time.time()
train(epoch, model, criterion_xent, criterion_htri, optimizer,
trainloader, use_gpu)
train_time += round(time.time() - start_train_time)
scheduler.step()
if (epoch + 1) > args.start_eval and args.eval_step > 0 and (
epoch + 1) % args.eval_step == 0 or (epoch +
1) == args.max_epoch:
print("==> Test")
rank1 = test(model, queryloader, galleryloader, use_gpu)
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,
'rank1': rank1,
'epoch': epoch,
}, is_best,
osp.join(args.save_dir,
'checkpoint_ep' + str(epoch + 1) + '.pth.tar'))
print("==> Best Rank-1 {:.1%}, achieved at epoch {}".format(
best_rank1, best_epoch))
elapsed = round(time.time() - start_time)
elapsed = str(datetime.timedelta(seconds=elapsed))
train_time = str(datetime.timedelta(seconds=train_time))
print(
"Finished. Total elapsed time (h:m:s): {}. Training time (h:m:s): {}.".
format(elapsed, train_time))
def main():
global args
set_random_seed(args.seed)
if not args.use_avai_gpus:
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices
use_gpu = torch.cuda.is_available()
if args.use_cpu: use_gpu = False
log_name = 'log_test.txt' if args.evaluate else 'log_train.txt'
sys.stdout = Logger(osp.join(args.save_dir, log_name))
print('==========\nArgs:{}\n=========='.format(args))
if use_gpu:
print('Currently using GPU {}'.format(args.gpu_devices))
cudnn.benchmark = True
else:
print('Currently using CPU, however, GPU is highly recommended')
print('Initializing video data manager')
dm = VideoDataManager(use_gpu, **video_dataset_kwargs(args))
trainloader, testloader_dict = dm.return_dataloaders()
print('Initializing model: {}'.format(args.arch))
model = models.init_model(name=args.arch,
num_classes=dm.num_train_pids,
loss={'xent', 'htri'},
pretrained=not args.no_pretrained,
use_gpu=use_gpu)
print('Model size: {:.3f} M'.format(count_num_param(model)))
if args.load_weights and check_isfile(args.load_weights):
load_pretrained_weights(model, args.load_weights)
if args.resume and check_isfile(args.resume):
checkpoint = torch.load(args.resume)
model.load_state_dict(checkpoint['state_dict'])
args.start_epoch = checkpoint['epoch'] + 1
best_rank1 = checkpoint['rank1']
print('Loaded checkpoint from "{}"'.format(args.resume))
print('- start_epoch: {}\n- rank1: {}'.format(args.start_epoch,
best_rank1))
model = nn.DataParallel(model).cuda() if use_gpu else model
criterion = CrossEntropyLoss(num_classes=dm.num_train_pids,
use_gpu=use_gpu,
label_smooth=args.label_smooth)
criterion_htri = TripletLoss(margin=args.margin)
optimizer = init_optimizer(model, **optimizer_kwargs(args))
scheduler = init_lr_scheduler(optimizer, **lr_scheduler_kwargs(args))
if args.evaluate:
print('Evaluate only')
for name in args.target_names:
print('Evaluating {} ...'.format(name))
queryloader = testloader_dict[name]['query']
galleryloader = testloader_dict[name]['gallery']
distmat = test(model,
queryloader,
galleryloader,
args.pool_tracklet_features,
use_gpu,
return_distmat=True)
if args.visualize_ranks:
visualize_ranked_results(distmat,
dm.return_testdataset_by_name(name),
save_dir=osp.join(
args.save_dir, 'ranked_results',
name),
topk=20)
return
start_time = time.time()
ranklogger = RankLogger(args.source_names, args.target_names)
train_time = 0
print('=> Start training')
if args.fixbase_epoch > 0:
print(
'Train {} for {} epochs while keeping other layers frozen'.format(
args.open_layers, args.fixbase_epoch))
initial_optim_state = optimizer.state_dict()
for epoch in range(args.fixbase_epoch):
start_train_time = time.time()
train(epoch,
model,
criterion_xent,
criterion_htri,
optimizer,
trainloader,
use_gpu,
fixbase=True)
train_time += round(time.time() - start_train_time)
print('Done. All layers are open to train for {} epochs'.format(
args.max_epoch))
optimizer.load_state_dict(initial_optim_state)
for epoch in range(args.start_epoch, args.max_epoch):
start_train_time = time.time()
train(epoch, model, criterion_xent, criterion_htri, optimizer,
trainloader, use_gpu)
train_time += round(time.time() - start_train_time)
scheduler.step()
if (epoch + 1) > args.start_eval and args.eval_freq > 0 and (
epoch + 1) % args.eval_freq == 0 or (epoch +
1) == args.max_epoch:
print('=> Test')
for name in args.target_names:
print('Evaluating {} ...'.format(name))
queryloader = testloader_dict[name]['query']
galleryloader = testloader_dict[name]['gallery']
rank1 = test(model, queryloader, galleryloader,
args.pool_tracklet_features, use_gpu)
ranklogger.write(name, epoch + 1, rank1)
save_checkpoint(
{
'state_dict': model.state_dict(),
'rank1': rank1,
'epoch': epoch,
}, False,
osp.join(args.save_dir,
'checkpoint_ep' + str(epoch + 1) + '.pth.tar'))
elapsed = round(time.time() - start_time)
elapsed = str(datetime.timedelta(seconds=elapsed))
train_time = str(datetime.timedelta(seconds=train_time))
print(
'Finished. Total elapsed time (h:m:s): {}. Training time (h:m:s): {}.'.
format(elapsed, train_time))
ranklogger.show_summary()
def main():
global args
set_random_seed(args.seed)
if not args.use_avai_gpus:
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices
use_gpu = torch.cuda.is_available()
if args.use_cpu: use_gpu = False
log_name = 'log_test.txt' if args.evaluate else 'log_train.txt'
sys.stdout = Logger(osp.join(args.save_dir, log_name))
print("==========\nArgs:{}\n==========".format(args))
if use_gpu:
print("Currently using GPU {}".format(args.gpu_devices))
cudnn.benchmark = True
else:
print("Currently using CPU, however, GPU is highly recommended")
print("Initializing MultiScale data manager")
dm = ImageDataManager(use_gpu,
scales=[224, 160],
**image_dataset_kwargs(args))
trainloader, testloader_dict = dm.return_dataloaders()
# sys.exit(0)
print("Initializing model: {}".format(args.arch))
model = models.init_model(name=args.arch,
num_classes=dm.num_train_pids,
input_size=args.width,
loss={'xent', 'htri'},
use_gpu=use_gpu)
print("Model size: {:.3f} M".format(count_num_param(model)))
print(model)
# criterion = CrossEntropyLoss(num_classes=dm.num_train_pids, use_gpu=use_gpu, label_smooth=args.label_smooth)
criterion_xent = CrossEntropyLoss(num_classes=dm.num_train_pids,
use_gpu=use_gpu,
label_smooth=args.label_smooth)
criterion_htri = TripletLoss(margin=args.margin)
optimizer = init_optimizer(model.parameters(), **optimizer_kwargs(args))
scheduler = lr_scheduler.MultiStepLR(optimizer,
milestones=args.stepsize,
gamma=args.gamma)
# scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.1, patience=3, verbose=True, threshold=1e-04)
if args.load_weights and check_isfile(
args.load_weights
): # load pretrained weights but ignore layers that don't match in size
checkpoint = torch.load(args.load_weights)
pretrain_dict = checkpoint['state_dict']
model_dict = model.state_dict()
pretrain_dict = {
k: v
for k, v in pretrain_dict.items()
if k in model_dict and model_dict[k].size() == v.size()
}
model_dict.update(pretrain_dict)
model.load_state_dict(model_dict)
print("Loaded pretrained weights from '{}'".format(args.load_weights))
if args.resume and check_isfile(args.resume):
checkpoint = torch.load(args.resume)
model.load_state_dict(checkpoint['state_dict'])
args.start_epoch = checkpoint['epoch'] + 1
print("Loaded checkpoint from '{}'".format(args.resume))
print("- start_epoch: {}\n- rank1: {}".format(args.start_epoch,
checkpoint['rank1']))
if use_gpu:
model = nn.DataParallel(model).cuda()
if args.evaluate:
print("Evaluate only")
for name in args.target_names:
print("Evaluating {} ...".format(name))
queryloader = testloader_dict[name]['query']
galleryloader = testloader_dict[name]['gallery']
test_set = dm.return_testdataset_by_name(name)
rank1, mAP = test(model,
test_set,
name,
queryloader,
galleryloader,
use_gpu,
visualize=args.visualize_ranks)
return
start_time = time.time()
ranklogger = RankLogger(args.source_names, args.target_names)
maplogger = RankLogger(args.source_names, args.target_names)
train_time = 0
# Tensorboard
writer = SummaryWriter(log_dir=osp.join('runs', 'tensorboard'))
print("=> Start training")
if args.fixbase_epoch > 0:
print(
"Train {} for {} epochs while keeping other layers frozen".format(
args.open_layers, args.fixbase_epoch))
initial_optim_state = optimizer.state_dict()
for epoch in range(args.fixbase_epoch):
start_train_time = time.time()
loss, prec1 = train(epoch,
model,
criterion,
optimizer,
trainloader,
writer,
use_gpu,
fixbase=True)
writer.add_scalar('train/loss', loss, epoch + 1)
writer.add_scalar('train/prec1', prec1, epoch + 1)
print(
'Epoch: [{:02d}] [Average Loss:] {:.4f}\t [Average Prec.:] {:.2%}'
.format(epoch + 1, loss, prec1))
train_time += round(time.time() - start_train_time)
print("Done. All layers are open to train for {} epochs".format(
args.max_epoch))
optimizer.load_state_dict(initial_optim_state)
args.start_epoch += args.fixbase_epoch
args.max_epoch += args.fixbase_epoch
for epoch in range(args.start_epoch, args.max_epoch):
start_train_time = time.time()
loss, prec1 = train(epoch, model, criterion, optimizer, trainloader,
writer, use_gpu)
writer.add_scalar('train/loss', loss, epoch + 1)
writer.add_scalar('train/prec1', prec1, epoch + 1)
print(
'Epoch: [{:02d}] [Average Loss:] {:.4f}\t [Average Prec.:] {:.2%}'.
format(epoch + 1, loss, prec1))
train_time += round(time.time() - start_train_time)
scheduler.step()
if (epoch + 1) > args.start_eval and args.eval_freq > 0 and (
epoch + 1) % args.eval_freq == 0 or (epoch +
1) == args.max_epoch:
print("=> Test")
for name in args.target_names:
print("Evaluating {} ...".format(name))
queryloader = testloader_dict[name]['query']
galleryloader = testloader_dict[name]['gallery']
test_set = dm.return_testdataset_by_name(name)
if epoch + 1 == args.max_epoch:
rank1, mAP = test(model,
test_set,
name,
queryloader,
galleryloader,
use_gpu,
visualize=True)
else:
rank1, mAP = test(model, test_set, name, queryloader,
galleryloader, use_gpu)
writer.add_scalar('test/top1', rank1, epoch + 1)
writer.add_scalar('test/mAP', mAP, epoch + 1)
ranklogger.write(name, epoch + 1, rank1)
maplogger.write(name, epoch + 1, mAP)
if use_gpu:
state_dict = model.module.state_dict()
else:
state_dict = model.state_dict()
save_checkpoint(
{
'state_dict': state_dict,
'rank1': rank1,
'epoch': epoch,
}, False,
osp.join(args.save_dir,
'checkpoint_ep' + str(epoch + 1) + '.pth.tar'))
# save last checkpoint
save_checkpoint(
{
'state_dict': state_dict,
'rank1': rank1,
'epoch': epoch,
}, False,
osp.join(args.save_dir, 'checkpoint_ep' + str(epoch + 1) + '.pth.tar'))
elapsed = round(time.time() - start_time)
elapsed = str(datetime.timedelta(seconds=elapsed))
train_time = str(datetime.timedelta(seconds=train_time))
print(
"Finished. Total elapsed time (h:m:s): {}. Training time (h:m:s): {}.".
format(elapsed, train_time))
ranklogger.show_summary()
maplogger.show_summary()