def main(args):
path = os.path.join(os.getcwd(), 'soft_label', 'soft_label_resnet50.txt')
if not os.path.isfile(path):
print('soft label file is not exist')
train_loader = getTrainLoader(args, path)
_, val_loader, num_query, num_classes, train_size = make_data_loader(args)
#train_loader, val_loader, num_query, num_classes, train_size = make_data_loader(args)
model = build_model(args, num_classes)
optimizer = make_optimizer(args, model)
scheduler = WarmupMultiStepLR(optimizer, [30, 55], 0.1, 0.01, 5, "linear")
loss_func = make_loss(args)
model.to(device)
for epoch in range(args.Epochs):
model.train()
running_loss = 0.0
running_klloss = 0.0
running_softloss = 0.0
running_corrects = 0.0
for index, data in enumerate(tqdm(train_loader)):
img, target, soft_target = data
img = img.cuda()
target = target.cuda()
soft_target = soft_target.cuda()
score, _ = model(img)
preds = torch.max(score.data, 1)[1]
loss, klloss, softloss = loss_func(score, target, soft_target)
optimizer.zero_grad()
loss.backward()
optimizer.step()
running_loss += loss.item()
running_klloss += klloss.item()
running_softloss += softloss.item()
running_corrects += float(torch.sum(preds == target.data))
scheduler.step()
epoch_loss = running_loss / train_size
epoch_klloss = running_klloss / train_size
epoch_softloss = running_softloss / train_size
epoch_acc = running_corrects / train_size
print(
"Epoch {} Loss : {:.4f} KLLoss:{:.8f} SoftLoss:{:.4f} Acc:{:.4f}"
.format(epoch, epoch_loss, epoch_klloss, epoch_softloss,
epoch_acc))
if (epoch + 1) % args.n_save == 0:
evaluator = Evaluator(model, val_loader, num_query)
cmc, mAP = evaluator.run()
print('---------------------------')
print("CMC Curve:")
for r in [1, 5, 10]:
print("Rank-{} : {:.1%}".format(r, cmc[r - 1]))
print("mAP : {:.1%}".format(mAP))
print('---------------------------')
save_model(args, model, optimizer, epoch)
def train(cfg):
logger = setup_logger("reid_baseline", cfg.OUTPUT_DIR)
logger.info("Running with config:\n{}".format(cfg))
# prepare camstyle dataset
train_loader, train_camstyle_loader, val_loader, num_query, num_classes = make_camstyle_data_loader(
cfg)
num_classes.append(-1)
# prepare model
model = build_model(cfg, num_classes)
optimizer, _ = make_optimizer(cfg, model)
loss_fn = make_loss(cfg, num_classes)
# Add for using self trained model
if cfg.MODEL.PRETRAIN_CHOICE == 'resume':
start_epoch = eval(
cfg.MODEL.PRETRAIN_PATH.split('/')[-1].split('.')[0].split('_')
[-1])
logger.info('Start epoch:%d' % start_epoch)
path_to_optimizer = cfg.MODEL.PRETRAIN_PATH.replace(
'model', 'optimizer')
logger.info('Path to the checkpoint of optimizer:%s' %
path_to_optimizer)
model.load_state_dict(torch.load(cfg.MODEL.PRETRAIN_PATH))
optimizer.load_state_dict(torch.load(path_to_optimizer))
scheduler = WarmupMultiStepLR(optimizer, cfg.SOLVER.STEPS,
cfg.SOLVER.GAMMA,
cfg.SOLVER.WARMUP_FACTOR,
cfg.SOLVER.WARMUP_ITERS,
cfg.SOLVER.WARMUP_METHOD, start_epoch)
elif cfg.MODEL.PRETRAIN_CHOICE == 'self' or cfg.MODEL.PRETRAIN_CHOICE == 'imagenet':
start_epoch = 0
model.load_param(cfg.MODEL.PRETRAIN_PATH, cfg.MODEL.PRETRAIN_CHOICE)
scheduler = WarmupMultiStepLR(optimizer, cfg.SOLVER.STEPS,
cfg.SOLVER.GAMMA,
cfg.SOLVER.WARMUP_FACTOR,
cfg.SOLVER.WARMUP_ITERS,
cfg.SOLVER.WARMUP_METHOD)
else:
logger.info(
'Only support pretrain_choice for imagenet and self, but got {}'.
format(cfg.MODEL.PRETRAIN_CHOICE))
do_train(
cfg,
model,
train_loader,
val_loader,
optimizer,
scheduler, # modify for using self trained model
loss_fn,
num_query,
start_epoch, # add for using self trained model
0,
train_camstyle_loader)
def _init_optimizers(self):
self.Content_optimizer = make_optimizer(self.cfg, self.Content_Encoder)
self.Content_optimizer_fix = make_optimizer(self.cfg, self.Content_Encoder, fix=True)
self.scheduler = WarmupMultiStepLR(self.Content_optimizer, (30, 55), 0.1, 1.0 / 3,
500, "linear")
self.scheduler_fix = WarmupMultiStepLR(self.Content_optimizer_fix, (30, 55), 0.1, 1.0 / 3,
500, "linear")
self.schedulers = []
self.optimizers = []
def train(cfg):
# prepare dataset
train_loader, val_loader, num_query, num_classes = make_data_loader(cfg)
# prepare model
model = build_model(cfg, num_classes)
total = sum([param.nelement() for param in model.parameters()])
print("Number of parameter: %.2fM" % (total / 1e6))
if cfg.MODEL.METRIC_LOSS_TYPE == 'triplet':
print('Train without center loss, the loss type is',
cfg.MODEL.METRIC_LOSS_TYPE)
optimizer = make_optimizer(cfg, model)
loss_func = make_loss(cfg, num_classes)
# Add for using self trained model
if cfg.MODEL.PRETRAIN_CHOICE == 'self':
start_epoch = eval(
cfg.MODEL.PRETRAIN_PATH.split('/')[-1].split('.')[0].split('_')
[-1])
print('Start epoch:', start_epoch)
path_to_optimizer = cfg.MODEL.PRETRAIN_PATH.replace(
'model', 'optimizer')
print('Path to the checkpoint of optimizer:', path_to_optimizer)
model.load_state_dict(torch.load(cfg.MODEL.PRETRAIN_PATH))
optimizer.load_state_dict(torch.load(path_to_optimizer))
scheduler = WarmupMultiStepLR(optimizer, cfg.SOLVER.STEPS,
cfg.SOLVER.GAMMA,
cfg.SOLVER.WARMUP_FACTOR,
cfg.SOLVER.WARMUP_ITERS,
cfg.SOLVER.WARMUP_METHOD,
start_epoch)
elif cfg.MODEL.PRETRAIN_CHOICE == 'imagenet':
start_epoch = 0
scheduler = WarmupMultiStepLR(optimizer, cfg.SOLVER.STEPS,
cfg.SOLVER.GAMMA,
cfg.SOLVER.WARMUP_FACTOR,
cfg.SOLVER.WARMUP_ITERS,
cfg.SOLVER.WARMUP_METHOD)
else:
print(
'Only support pretrain_choice for imagenet and self, but got {}'
.format(cfg.MODEL.PRETRAIN_CHOICE))
do_train(
cfg,
model,
train_loader,
val_loader,
optimizer,
scheduler, # modify for using self trained model
loss_func,
num_query,
start_epoch # add for using self trained model
)
def main(args):
sys.stdout = Logger(
os.path.join(args.log_path, args.log_description,
'log' + time.strftime(".%m_%d_%H:%M:%S") + '.txt'))
train_loader, val_loader, num_query, num_classes, train_size = make_data_loader(
args)
model = build_model(args, num_classes)
print(model)
optimizer = make_optimizer(args, model)
scheduler = WarmupMultiStepLR(optimizer, [30, 55], 0.1, 0.01, 5, "linear")
loss_func = make_loss(args)
model.to(device)
for epoch in range(args.Epochs):
model.train()
running_loss = 0.0
running_corrects = 0.0
for index, data in enumerate(tqdm(train_loader)):
img, target = data
img = img.cuda()
target = target.cuda()
score, _ = model(img)
preds = torch.max(score.data, 1)[1]
loss = loss_func(score, target)
optimizer.zero_grad()
loss.backward()
optimizer.step()
running_loss += loss.item()
running_corrects += float(torch.sum(preds == target.data))
scheduler.step()
epoch_loss = running_loss / train_size
epoch_acc = running_corrects / train_size
print("Epoch {} Loss : {:.6f} Acc:{:.4f}".format(
epoch, epoch_loss, epoch_acc))
if (epoch + 1) % args.n_save == 0:
evaluator = Evaluator(model, val_loader, num_query)
cmc, mAP = evaluator.run()
print('---------------------------')
print("CMC Curve:")
for r in [1, 5, 10]:
print("Rank-{} : {:.1%}".format(r, cmc[r - 1]))
print("mAP : {:.1%}".format(mAP))
print('---------------------------')
save_model(args, model, optimizer, epoch)
def train(cfg, cfg_hr):
# prepare dataset
train_loader, val_loader, num_query, num_classes = make_data_loader(cfg)
# prepare model
model = build_model(cfg, cfg_hr, num_classes)
model = nn.DataParallel(model)
if cfg.MODEL.IF_WITH_CENTER == 'no':
print('Train without center loss, the loss type is', cfg.MODEL.METRIC_LOSS_TYPE)
optimizer = make_optimizer(cfg, model)
print(cfg.SOLVER.MARGIN)
scheduler = WarmupMultiStepLR(optimizer, cfg.SOLVER.STEPS, cfg.SOLVER.GAMMA, cfg.SOLVER.WARMUP_FACTOR,
cfg.SOLVER.WARMUP_ITERS, cfg.SOLVER.WARMUP_METHOD)
loss_func = make_loss(cfg, num_classes) # modified by gu
arguments = {}
do_train(
cfg,
model,
train_loader,
val_loader,
optimizer,
scheduler,
loss_func,
num_query
)
elif cfg.MODEL.IF_WITH_CENTER == 'yes':
print('Train with center loss, the loss type is', cfg.MODEL.METRIC_LOSS_TYPE)
loss_func, center_criterion = make_loss_with_center(cfg, num_classes) # modified by gu
optimizer, optimizer_center = make_optimizer_with_center(cfg, model, center_criterion)
scheduler = WarmupMultiStepLR(optimizer, cfg.SOLVER.STEPS, cfg.SOLVER.GAMMA, cfg.SOLVER.WARMUP_FACTOR,
cfg.SOLVER.WARMUP_ITERS, cfg.SOLVER.WARMUP_METHOD)
arguments = {}
do_train_with_center(
cfg,
model,
center_criterion,
train_loader,
val_loader,
optimizer,
optimizer_center,
scheduler,
loss_func,
num_query
)
else:
print("Unsupported value for cfg.MODEL.IF_WITH_CENTER {}, only support yes or no!\n".format(cfg.MODEL.IF_WITH_CENTER))
def train(cfg):
logger = setup_logger("reid_baseline", cfg.OUTPUT_DIR)
logger.info("Running with config:\n{}".format(cfg))
# prepare dataset
val_data_loader, num_query = make_val_data_loader(cfg)
num_classes = np.zeros(len(cfg.DATALOADER.SAMPLER_PROB)).astype(int) - 1
source_dataset = init_dataset(cfg.SRC_DATA.NAMES, root_train=cfg.SRC_DATA.TRAIN_DIR, transfered=cfg.SRC_DATA.TRANSFERED)
num_classes[0] = source_dataset.num_train_pids
num_classes[1] = cfg.TGT_UNSUPDATA.CLUSTER_TOPK
if cfg.MODEL.FINETUNE:
num_classes[1] += 200
# prepare model
model = build_model(cfg, num_classes)
optimizer,fixed_lr_idxs = make_optimizer(cfg, model)
loss_fn = make_loss(cfg, num_classes)
# Add for using self trained model
if cfg.MODEL.PRETRAIN_CHOICE == 'resume':
start_epoch = eval(cfg.MODEL.PRETRAIN_PATH.split('/')[-1].split('.')[0].split('_')[-1])
logger.info('Start epoch:%d' %start_epoch)
path_to_optimizer = cfg.MODEL.PRETRAIN_PATH.replace('model', 'optimizer')
logger.info('Path to the checkpoint of optimizer:%s' %path_to_optimizer)
model.load_state_dict(torch.load(cfg.MODEL.PRETRAIN_PATH))
optimizer.load_state_dict(torch.load(path_to_optimizer))
scheduler = WarmupMultiStepLR(optimizer, cfg.SOLVER.STEPS, cfg.SOLVER.GAMMA, cfg.SOLVER.WARMUP_FACTOR, cfg.SOLVER.WARMUP_ITERS, cfg.SOLVER.WARMUP_METHOD, start_epoch, fixed_lr_idxs)
elif cfg.MODEL.PRETRAIN_CHOICE == 'self' or cfg.MODEL.PRETRAIN_CHOICE == 'imagenet':
start_epoch = 0
model.load_param(cfg.MODEL.PRETRAIN_PATH, cfg.MODEL.PRETRAIN_CHOICE)
scheduler = WarmupMultiStepLR(optimizer, cfg.SOLVER.STEPS, cfg.SOLVER.GAMMA, cfg.SOLVER.WARMUP_FACTOR, cfg.SOLVER.WARMUP_ITERS, cfg.SOLVER.WARMUP_METHOD, -1, fixed_lr_idxs)
camera_model = build_camera_model(cfg, num_classes=5)
camera_model.load_param(cfg.TEST.CAMERA_WEIGHT, cfg.MODEL.PRETRAIN_CHOICE)
else:
logger.info('Only support pretrain_choice for imagenet and self, but got {}'.format(cfg.MODEL.PRETRAIN_CHOICE))
do_train(cfg,
model,
camera_model,
val_data_loader,
optimizer,
scheduler, # modify for using self trained model
loss_fn,
num_query,
start_epoch, # add for using self trained model
0
)
def train(cfg):
# prepare dataset
train_loader, val_loader, num_query, num_classes = make_data_loader(cfg)
# prepare model
model = build_model(cfg, num_classes)
print('Train without center loss, the loss type is',
cfg.MODEL.METRIC_LOSS_TYPE)
optimizer = make_optimizer(cfg, model)
loss_func = make_loss(cfg, num_classes)
if cfg.MODEL.PRETRAIN_CHOICE == 'self':
start_epoch = eval(
cfg.MODEL.PRETRAIN_PATH.split('/')[-1].split('.')[0].split('_')
[-1])
print('Start epoch:', start_epoch)
path_to_optimizer = cfg.MODEL.PRETRAIN_PATH.replace(
'model', 'optimizer')
print('Path to the checkpoint of optimizer:', path_to_optimizer)
model.load_state_dict(torch.load(cfg.MODEL.PRETRAIN_PATH))
optimizer.load_state_dict(torch.load(path_to_optimizer))
if cfg.MODEL.DEVICE == "cuda":
for state in optimizer.state.values():
for k, v in state.items():
if isinstance(v, torch.Tensor):
state[k] = v.cuda()
scheduler = WarmupMultiStepLR(optimizer, cfg.SOLVER.STEPS,
cfg.SOLVER.GAMMA,
cfg.SOLVER.WARMUP_FACTOR,
cfg.SOLVER.WARMUP_ITERS,
cfg.SOLVER.WARMUP_METHOD, start_epoch)
elif cfg.MODEL.PRETRAIN_CHOICE == 'imagenet':
start_epoch = 0
scheduler = WarmupMultiStepLR(optimizer, cfg.SOLVER.STEPS,
cfg.SOLVER.GAMMA,
cfg.SOLVER.WARMUP_FACTOR,
cfg.SOLVER.WARMUP_ITERS,
cfg.SOLVER.WARMUP_METHOD)
else:
start_epoch = 0
print('Only support pretrain_choice for imagenet and self, but got {}'.
format(cfg.MODEL.PRETRAIN_CHOICE))
arguments = {}
do_train(cfg, model, train_loader, val_loader, optimizer, scheduler,
loss_func, num_query, start_epoch)
def __init__(self, cfg, num_classes):
self.device = cfg.MODEL.DEVICE
self.model = build_model(cfg, num_classes)
self.loss = Loss(cfg, num_classes, self.model.in_planes)
self.optimizer = make_optimizer(cfg, self.model)
self.scheduler = WarmupMultiStepLR(self.optimizer,
cfg.WARMUP.STEPS,
cfg.WARMUP.GAMMA,
cfg.WARMUP.FACTOR,
cfg.WARMUP.MAX_EPOCHS,
cfg.WARMUP.METHOD)
if cfg.APEX.IF_ON:
logger.info("Using apex")
try:
import apex
except ImportError:
raise ImportError("Please install apex from https://www.github.com/nvidia/apex to run this example.")
assert torch.backends.cudnn.enabled, "Amp requires cudnn backend to be enabled."
#
# if cfg.APEX.IF_SYNC_BN:
# logger.info("Using apex synced BN")
# self.module = apex.parallel.convert_syncbn_model(self.module)
if self.device is 'cuda':
self.model = self.model.cuda()
if cfg.APEX.IF_ON:
from apex import amp
self.model, self.optimizer = amp.initialize(self.model,
self.optimizer,
opt_level=cfg.APEX.OPT_LEVEL,
keep_batchnorm_fp32=None if cfg.APEX.OPT_LEVEL == 'O1' else True,
loss_scale=cfg.APEX.LOSS_SCALE[0])
def train(cfg):
# prepare dataset
train_loader, val_loader, num_query, num_classes = make_data_loader(cfg)
# prepare model
model = build_model(cfg, num_classes)
if cfg.SOLVER.FINETUNE:
model.load_state_dict(torch.load(cfg.TEST.WEIGHT).module.state_dict())
model = nn.DataParallel(model)
optimizer = make_optimizer(cfg, model)
scheduler = WarmupMultiStepLR(optimizer, cfg.SOLVER.STEPS, cfg.SOLVER.GAMMA, cfg.SOLVER.WARMUP_FACTOR,
cfg.SOLVER.WARMUP_ITERS, cfg.SOLVER.WARMUP_METHOD)
# scheduler = WarmupStepLR(optimizer,3, 9, cfg.SOLVER.WARMUP_FACTOR,
# cfg.SOLVER.WARMUP_ITERS, cfg.SOLVER.WARMUP_METHOD)
loss_func = make_loss(cfg)
arguments = {}
do_train(
cfg,
model,
train_loader,
val_loader,
optimizer,
scheduler,
loss_func,
num_query
)
def main():
output_dir = cfg.MODEL.OUTPUT_DIR
if output_dir and not os.path.exists(output_dir):
os.makedirs(output_dir)
logger = setup_logger("tracker", output_dir, 0)
logger.info("Running with config:\n{}".format(cfg))
torch.backends.cudnn.benchmark = True
train_loader, val_loader = make_data_loader(cfg)
model = build_model(cfg)
optimizer = make_optimizer(cfg, model)
scheduler = WarmupMultiStepLR(optimizer, cfg.SOLVER.STEPS, cfg.SOLVER.GAMMA, cfg.SOLVER.WARMUP_FACTOR,
cfg.SOLVER.WARMUP_ITERS, cfg.SOLVER.WARMUP_METHOD)
loss_func = make_loss(cfg)
do_train(
cfg,
model,
train_loader,
val_loader,
optimizer,
scheduler,
loss_func
)
def train(cfg):
# prepare dataset
train_loader, val_loader, num_query, num_classes, clustering_loader = make_data_loader(
cfg)
# prepare model
model = build_model(cfg, num_classes)
if cfg.MODEL.IF_WITH_CENTER == 'on':
loss_func, center_criterion_part, center_criterion_global, center_criterion_fore = make_loss_with_center(
cfg, num_classes)
optimizer, optimizer_center = make_optimizer_with_center(
cfg, model, center_criterion_part, center_criterion_global,
center_criterion_fore)
else:
loss_func = make_loss(cfg, num_classes)
optimizer = make_optimizer(cfg, model)
# Add for using self trained model
if cfg.MODEL.PRETRAIN_CHOICE == 'imagenet':
start_epoch = 0
scheduler = WarmupMultiStepLR(optimizer, cfg.SOLVER.STEPS,
cfg.SOLVER.GAMMA,
cfg.SOLVER.WARMUP_FACTOR,
cfg.SOLVER.WARMUP_ITERS,
cfg.SOLVER.WARMUP_METHOD)
else:
print('Only support pretrain_choice for imagenet, but got {}'.format(
cfg.MODEL.PRETRAIN_CHOICE))
if cfg.MODEL.IF_WITH_CENTER == 'on':
do_train_with_center(
cfg,
model,
center_criterion_part,
center_criterion_global,
center_criterion_fore,
train_loader,
val_loader,
optimizer,
optimizer_center,
scheduler, # modify for using self trained model
loss_func,
num_query,
start_epoch, # add for using self trained model
clustering_loader)
else:
do_train(
cfg,
model,
train_loader,
val_loader,
optimizer,
scheduler, # modify for using self trained model
loss_func,
num_query,
start_epoch, # add for using self trained model
clustering_loader)
def train(cfg):
# prepare dataset
train_loader, val_loader, num_query, num_classes, dataset = make_data_loader(cfg)
# prepare model
model = build_model(cfg, num_classes)
print('Train without center loss, the loss type is', cfg['MODEL.METRIC_LOSS_TYPE'])
optimizer = make_optimizer(cfg, model)
loss_func = make_loss(cfg, num_classes, model.in_planes) # modified by gu
# Add for using self trained model
if cfg['MODEL.PRETRAIN_CHOICE'] == 'continue':
#start_epoch = eval(cfg['MODEL.PRETRAIN_PATH'].split('/')[-1].split('.')[0].split('_')[-1])
start_epoch = 0
print('Start epoch:', start_epoch)
path_to_optimizer = cfg['MODEL.PRETRAIN_PATH'].replace('model', 'optimizer')
print('Path to the checkpoint of optimizer:', path_to_optimizer)
model.load_param(cfg['MODEL.PRETRAIN_PATH'])
#optimizer = make_optimizer(cfg, model)
scheduler = WarmupMultiStepLR(optimizer, cfg['SOLVER.STEPS'], cfg['SOLVER.GAMMA'], cfg['SOLVER.WARMUP_FACTOR'],
cfg['SOLVER.WARMUP_ITERS'], cfg['SOLVER.WARMUP_METHOD'])
elif cfg['MODEL.PRETRAIN_CHOICE'] == 'imagenet':
start_epoch = 0
scheduler = WarmupMultiStepLR(optimizer, cfg['SOLVER.STEPS'], cfg['SOLVER.GAMMA'], cfg['SOLVER.WARMUP_FACTOR'],
cfg['SOLVER.WARMUP_ITERS'], cfg['SOLVER.WARMUP_METHOD'])
elif cfg['MODEL.PRETRAIN_CHOICE'] == 'self' or cfg['MODEL.PRETRAIN_CHOICE'] == 'self-no-head':
start_epoch = 0
scheduler = WarmupMultiStepLR(optimizer, cfg['SOLVER.STEPS'], cfg['SOLVER.GAMMA'],
cfg['SOLVER.WARMUP_FACTOR'],
cfg['SOLVER.WARMUP_ITERS'], cfg['SOLVER.WARMUP_METHOD'])
else:
print('Only support pretrain_choice for imagenet and self, but got {}'.format(cfg['MODEL.PRETRAIN_CHOICE']))
do_train(
cfg,
model,
train_loader,
val_loader,
optimizer,
scheduler, # modify for using self trained model
loss_func,
num_query,
start_epoch, # add for using self trained model
dataset
)
def get_component(self, cfg, num_classes):
self.device = cfg.MODEL.DEVICE
self.model = build_model(cfg, num_classes)
self.loss = Loss(cfg, num_classes, self.model.in_planes)
self.optimizer = make_optimizer(cfg, self.model)
self.scheduler = WarmupMultiStepLR(self.optimizer,
cfg.WARMUP.STEPS,
cfg.WARMUP.GAMMA,
cfg.WARMUP.FACTOR,
cfg.WARMUP.MAX_EPOCHS,
cfg.WARMUP.METHOD)
def train(cfg):
# prepare dataset
train_loader, val_loader, num_query, num_classes = make_data_loader(cfg)
# prepare model
model = build_model(cfg, num_classes)
optimizer = make_optimizer(cfg, model)
scheduler = WarmupMultiStepLR(optimizer, cfg.SOLVER.STEPS,
cfg.SOLVER.GAMMA, cfg.SOLVER.WARMUP_FACTOR,
cfg.SOLVER.WARMUP_ITERS,
cfg.SOLVER.WARMUP_METHOD)
loss_func = make_loss(cfg)
arguments = {}
do_train(cfg, model, train_loader, val_loader, optimizer, scheduler,
loss_func, num_query)
if __name__ == '__main__':
Cfg = Config()
logger = setup_logger('{}'.format(Cfg.PROJECT_NAME), Cfg.LOG_DIR)
logger.info("Running with config:\n{}".format(Cfg.PROJECT_NAME))
os.environ['CUDA_VISIBLE_DEVICES'] = Cfg.DEVICE_ID
cudnn.benchmark = True
# This flag allows you to enable the inbuilt cudnn auto-tuner to find the best algorithm to use for your hardware.
train_loader, val_loader, num_query, num_classes = make_dataloader(Cfg)
model = make_model(Cfg, num_class=num_classes)
loss_func, center_criterion = make_loss(Cfg, num_classes=num_classes)
optimizer, optimizer_center = make_optimizer(Cfg, model, center_criterion)
scheduler = WarmupMultiStepLR(optimizer, Cfg.STEPS, Cfg.GAMMA,
Cfg.WARMUP_FACTOR, Cfg.WARMUP_EPOCHS,
Cfg.WARMUP_METHOD)
do_train(
Cfg,
model,
center_criterion,
train_loader,
val_loader,
optimizer,
optimizer_center,
scheduler, # modify for using self trained model
loss_func,
num_query)
logger = setup_logger('pose-transfer-gan.train', log_dir)
logger.info("Running with config:\n{}".format(Cfg))
os.environ['CUDA_VISIBLE_DEVICES'] = Cfg.MODEL.DEVICE_ID
cudnn.benchmark = True
# This flag allows you to enable the inbuilt cudnn auto-tuner to find the best algorithm to use for your hardware.
train_loader, val_loader = make_dataloader(Cfg)
model_G, model_Dip, model_Dii, model_D_reid = make_model(Cfg)
optimizerG = make_optimizer(Cfg, model_G)
optimizerDip = make_optimizer(Cfg, model_Dip)
optimizerDii = make_optimizer(Cfg, model_Dii)
schedulerG = WarmupMultiStepLR(optimizerG, Cfg.SOLVER.STEPS,
Cfg.SOLVER.GAMMA, Cfg.SOLVER.WARMUP_FACTOR,
Cfg.SOLVER.WARMUP_EPOCHS,
Cfg.SOLVER.WARMUP_METHOD)
schedulerDip = WarmupMultiStepLR(optimizerDip, Cfg.SOLVER.STEPS,
Cfg.SOLVER.GAMMA,
Cfg.SOLVER.WARMUP_FACTOR,
Cfg.SOLVER.WARMUP_EPOCHS,
Cfg.SOLVER.WARMUP_METHOD)
schedulerDii = WarmupMultiStepLR(optimizerDii, Cfg.SOLVER.STEPS,
Cfg.SOLVER.GAMMA,
Cfg.SOLVER.WARMUP_FACTOR,
Cfg.SOLVER.WARMUP_EPOCHS,
Cfg.SOLVER.WARMUP_METHOD)
GAN_loss, L1_loss, ReID_loss = make_loss(Cfg)
do_train(Cfg, model_G, model_Dip, model_Dii, model_D_reid, train_loader,
val_loader, optimizerG, optimizerDip, optimizerDii, GAN_loss,
L1_loss, ReID_loss, schedulerG, schedulerDip, schedulerDii)
def train(args):
if args.batch_size % args.num_instance != 0:
new_batch_size = (args.batch_size //
args.num_instance) * args.num_instance
print(
f"given batch size is {args.batch_size} and num_instances is {args.num_instance}."
+
f"Batch size must be divided into {args.num_instance}. Batch size will be replaced into {new_batch_size}"
)
args.batch_size = new_batch_size
# prepare dataset
train_loader, val_loader, num_query, train_data_len, num_classes = make_data_loader(
args)
model = build_model(args, num_classes)
print("model size: {:.5f}M".format(
sum(p.numel() for p in model.parameters()) / 1e6))
loss_fn, center_criterion = make_loss(args, num_classes)
optimizer, optimizer_center = make_optimizer(args, model, center_criterion)
if args.cuda:
model = model.cuda()
if args.amp:
if args.center_loss:
model, [optimizer, optimizer_center] = \
amp.initialize(model, [optimizer, optimizer_center], opt_level="O1")
else:
model, optimizer = amp.initialize(model,
optimizer,
opt_level="O1")
for state in optimizer.state.values():
for k, v in state.items():
if isinstance(v, torch.Tensor):
state[k] = v.cuda()
if args.center_loss:
center_criterion = center_criterion.cuda()
for state in optimizer_center.state.values():
for k, v in state.items():
if isinstance(v, torch.Tensor):
state[k] = v.cuda()
model_state_dict = model.state_dict()
optim_state_dict = optimizer.state_dict()
if args.center_loss:
optim_center_state_dict = optimizer_center.state_dict()
center_state_dict = center_criterion.state_dict()
reid_evaluator = ReIDEvaluator(args, model, num_query)
start_epoch = 0
global_step = 0
if args.pretrain != '': # load pre-trained model
weights = torch.load(args.pretrain)
model_state_dict = weights["state_dict"]
model.load_state_dict(model_state_dict)
if args.center_loss:
center_criterion.load_state_dict(
torch.load(args.pretrain.replace(
'model', 'center_param'))["state_dict"])
if args.resume:
start_epoch = weights["epoch"]
global_step = weights["global_step"]
optimizer.load_state_dict(
torch.load(args.pretrain.replace('model',
'optimizer'))["state_dict"])
if args.center_loss:
optimizer_center.load_state_dict(
torch.load(
args.pretrain.replace(
'model', 'optimizer_center'))["state_dict"])
print(f'Start epoch: {start_epoch}, Start step: {global_step}')
scheduler = WarmupMultiStepLR(optimizer, args.steps, args.gamma,
args.warmup_factor, args.warmup_step,
"linear",
-1 if start_epoch == 0 else start_epoch)
current_epoch = start_epoch
best_epoch = 0
best_rank1 = 0
best_mAP = 0
if args.resume:
rank, mAP = reid_evaluator.evaluate(val_loader)
best_rank1 = rank[0]
best_mAP = mAP
best_epoch = current_epoch + 1
batch_time = AverageMeter()
total_losses = AverageMeter()
model_save_dir = os.path.join(args.save_dir, 'ckpts')
os.makedirs(model_save_dir, exist_ok=True)
summary_writer = SummaryWriter(log_dir=os.path.join(
args.save_dir, "tensorboard_log"),
purge_step=global_step)
def summary_loss(score, feat, labels, top_name='global'):
loss = 0.0
losses = loss_fn(score, feat, labels)
for loss_name, loss_val in losses.items():
if loss_name.lower() == "accuracy":
summary_writer.add_scalar(f"Score/{top_name}/triplet",
loss_val, global_step)
continue
if "dist" in loss_name.lower():
summary_writer.add_histogram(f"Distance/{loss_name}", loss_val,
global_step)
continue
loss += loss_val
summary_writer.add_scalar(f"losses/{top_name}/{loss_name}",
loss_val, global_step)
ohe_labels = torch.zeros_like(score)
ohe_labels.scatter_(1, labels.unsqueeze(1), 1.0)
cls_score = torch.softmax(score, dim=1)
cls_score = torch.sum(cls_score * ohe_labels, dim=1).mean()
summary_writer.add_scalar(f"Score/{top_name}/X-entropy", cls_score,
global_step)
return loss
def save_weights(file_name, eph, steps):
torch.save(
{
"state_dict": model_state_dict,
"epoch": eph + 1,
"global_step": steps
}, file_name)
torch.save({"state_dict": optim_state_dict},
file_name.replace("model", "optimizer"))
if args.center_loss:
torch.save({"state_dict": center_state_dict},
file_name.replace("model", "optimizer_center"))
torch.save({"state_dict": optim_center_state_dict},
file_name.replace("model", "center_param"))
# training start
for epoch in range(start_epoch, args.max_epoch):
model.train()
t0 = time.time()
for i, (inputs, labels, _, _) in enumerate(train_loader):
if args.cuda:
inputs = inputs.cuda()
labels = labels.cuda()
cls_scores, features = model(inputs, labels)
# losses
total_loss = summary_loss(cls_scores[0], features[0], labels,
'global')
if args.use_local_feat:
total_loss += summary_loss(cls_scores[1], features[1], labels,
'local')
optimizer.zero_grad()
if args.center_loss:
optimizer_center.zero_grad()
# backward with global loss
if args.amp:
optimizers = [optimizer]
if args.center_loss:
optimizers.append(optimizer_center)
with amp.scale_loss(total_loss, optimizers) as scaled_loss:
scaled_loss.backward()
else:
with torch.autograd.detect_anomaly():
total_loss.backward()
# optimization
optimizer.step()
if args.center_loss:
for name, param in center_criterion.named_parameters():
try:
param.grad.data *= (1. / args.center_loss_weight)
except AttributeError:
continue
optimizer_center.step()
batch_time.update(time.time() - t0)
total_losses.update(total_loss.item())
# learning_rate
current_lr = optimizer.param_groups[0]['lr']
summary_writer.add_scalar("lr", current_lr, global_step)
t0 = time.time()
if (i + 1) % args.log_period == 0:
print(
f"Epoch: [{epoch}][{i+1}/{train_data_len}] " +
f"Batch Time {batch_time.val:.3f} ({batch_time.mean:.3f}) "
+
f"Total_loss {total_losses.val:.3f} ({total_losses.mean:.3f})"
)
global_step += 1
print(
f"Epoch: [{epoch}]\tEpoch Time {batch_time.sum:.3f} s\tLoss {total_losses.mean:.3f}\tLr {current_lr:.2e}"
)
if args.eval_period > 0 and (epoch + 1) % args.eval_period == 0 or (
epoch + 1) == args.max_epoch:
rank, mAP = reid_evaluator.evaluate(
val_loader,
mode="retrieval" if args.dataset_name == "cub200" else "reid")
rank_string = ""
for r in (1, 2, 4, 5, 8, 10, 16, 20):
rank_string += f"Rank-{r:<3}: {rank[r-1]:.1%}"
if r != 20:
rank_string += " "
summary_writer.add_text("Recall@K", rank_string, global_step)
summary_writer.add_scalar("Rank-1", rank[0], (epoch + 1))
rank1 = rank[0]
is_best = rank1 > best_rank1
if is_best:
best_rank1 = rank1
best_mAP = mAP
best_epoch = epoch + 1
if (epoch + 1) % args.save_period == 0 or (epoch +
1) == args.max_epoch:
pth_file_name = os.path.join(
model_save_dir,
f"{args.backbone}_model_{epoch + 1}.pth.tar")
save_weights(pth_file_name, eph=epoch, steps=global_step)
if is_best:
pth_file_name = os.path.join(
model_save_dir, f"{args.backbone}_model_best.pth.tar")
save_weights(pth_file_name, eph=epoch, steps=global_step)
# end epoch
current_epoch += 1
batch_time.reset()
total_losses.reset()
torch.cuda.empty_cache()
# update learning rate
scheduler.step()
print(f"Best rank-1 {best_rank1:.1%}, achived at epoch {best_epoch}")
summary_writer.add_hparams(
{
"dataset_name": args.dataset_name,
"triplet_dim": args.triplet_dim,
"margin": args.margin,
"base_lr": args.base_lr,
"use_attn": args.use_attn,
"use_mask": args.use_mask,
"use_local_feat": args.use_local_feat
}, {
"mAP": best_mAP,
"Rank1": best_rank1
})
from loss import make_loss
from processor import do_train
from solver import make_optimizer, WarmupMultiStepLR
from utils.logger import setup_logger
if __name__ == '__main__':
Cfg = Configuration()
log_dir = Cfg.DATALOADER.LOG_DIR
logger = setup_logger('{}'.format(Cfg.PROJECT_NAME), log_dir)
logger.info("Running with config:\n{}".format(Cfg.PROJECT_NAME))
os.environ['CUDA_VISIBLE_DEVICES'] = Cfg.DEVICE_ID
cudnn.benchmark = True
# This flag allows you to enable the inbuilt cudnn auto-tuner to find the best algorithm to use for your hardware.
train_loader, val_loader = make_dataloader(Cfg)
model = make_model(Cfg)
optimizer = make_optimizer(Cfg, model)
scheduler = WarmupMultiStepLR(Cfg, optimizer)
loss_func = make_loss(Cfg)
do_train(
Cfg,
model,
train_loader,
val_loader,
optimizer,
scheduler, # modify for using self trained model
loss_func,
)
def train(cfg):
# prepare dataset
train_loader, val_loader, num_query, num_classes, attr_labels = make_mt_data_loader(
cfg)
print(f'attribute labels: {attr_labels}')
# prepare model
model = build_mt_model(num_features=cfg.mt.num_features,
dropout=cfg.mt.dropout,
last_stride=cfg.mt.last_conv_stride,
num_classes=num_classes,
num_classes_seg=cfg.mt.num_classes_seg,
global_branch=cfg.mt.global_branch,
mask_branch=cfg.mt.mask_branch,
part_branch=cfg.mt.part_branch,
mask_dim=cfg.mt.mask_dim,
part_dim=cfg.mt.part_dim,
part_info=cfg.mt.part_info,
attr_label_number=len(attr_labels),
attr_mask_weight=cfg.mt.attr_mask_weight,
wavp=cfg.mt.attr_wavp,
use_attr=cfg.mt.use_attr,
part_layer=cfg.mt.part_layer,
part_abla=cfg.mt.part_abla)
if cfg.MODEL.IF_WITH_CENTER == 'no':
print('Train without center loss, the loss type is',
cfg.MODEL.METRIC_LOSS_TYPE)
optimizer = make_optimizer(cfg, model)
# scheduler = WarmupMultiStepLR(optimizer, cfg.SOLVER.STEPS, cfg.SOLVER.GAMMA, cfg.SOLVER.WARMUP_FACTOR,
# cfg.SOLVER.WARMUP_ITERS, cfg.SOLVER.WARMUP_METHOD)
print('Builidng criterions')
loss_func = make_mt_loss(cfg, cfg.mt.normalize_size, cfg.mt.num_classes_seg, cfg.mt.weight[2],\
cfg.mt.triplet_margin, num_classes) # modified by gu
# Add for using self trained model
if cfg.MODEL.PRETRAIN_CHOICE == 'self':
start_epoch = eval(
cfg.MODEL.PRETRAIN_PATH.split('/')[-1].split('.')[0].split('_')
[-1])
print('Start epoch:', start_epoch)
path_to_optimizer = cfg.MODEL.PRETRAIN_PATH.replace(
'model', 'optimizer')
print('Path to the checkpoint of optimizer:', path_to_optimizer)
#try:
model.load_state_dict(torch.load(cfg.MODEL.PRETRAIN_PATH))
model.to(cfg.MODEL.DEVICE)
#def map_func(storage, location):
# return storage.cuda()
optimizer.load_state_dict(torch.load(path_to_optimizer))
#except:
# def map_func_cpu(storage, location):
# return storage.cpu()
# model.load_state_dict(torch.load(cfg.MODEL.PRETRAIN_PATH, map_location=map_func_cpu).state_dict())
# optimizer = torch.load(path_to_optimizer)
scheduler = WarmupMultiStepLR(optimizer, cfg.SOLVER.STEPS,
cfg.SOLVER.GAMMA,
cfg.SOLVER.WARMUP_FACTOR,
cfg.SOLVER.WARMUP_ITERS,
cfg.SOLVER.WARMUP_METHOD,
start_epoch)
elif cfg.MODEL.PRETRAIN_CHOICE == 'imagenet':
start_epoch = 0
scheduler = WarmupMultiStepLR(optimizer, cfg.SOLVER.STEPS,
cfg.SOLVER.GAMMA,
cfg.SOLVER.WARMUP_FACTOR,
cfg.SOLVER.WARMUP_ITERS,
cfg.SOLVER.WARMUP_METHOD)
else:
print(
'Only support pretrain_choice for imagenet and self, but got {}'
.format(cfg.MODEL.PRETRAIN_CHOICE))
arguments = {}
print('Runing Trainer...')
do_mt_train(
cfg,
model,
train_loader,
val_loader,
optimizer,
scheduler, # modify for using self trained model
loss_func,
num_query,
start_epoch, # add for using self trained model
cfg.mt.weight)
elif cfg.MODEL.IF_WITH_CENTER == 'yes':
print('Train with center loss, the loss type is',
cfg.MODEL.METRIC_LOSS_TYPE)
loss_func, center_criterion = make_loss_with_center(
cfg, num_classes) # modified by gu
optimizer, optimizer_center = make_optimizer_with_center(
cfg, model, center_criterion)
# scheduler = WarmupMultiStepLR(optimizer, cfg.SOLVER.STEPS, cfg.SOLVER.GAMMA, cfg.SOLVER.WARMUP_FACTOR,
# cfg.SOLVER.WARMUP_ITERS, cfg.SOLVER.WARMUP_METHOD)
arguments = {}
# Add for using self trained model
if cfg.MODEL.PRETRAIN_CHOICE == 'self':
start_epoch = eval(
cfg.MODEL.PRETRAIN_PATH.split('/')[-1].split('.')[0].split('_')
[-1])
print('Start epoch:', start_epoch)
path_to_optimizer = cfg.MODEL.PRETRAIN_PATH.replace(
'model', 'optimizer')
print('Path to the checkpoint of optimizer:', path_to_optimizer)
path_to_center_param = cfg.MODEL.PRETRAIN_PATH.replace(
'model', 'center_param')
print('Path to the checkpoint of center_param:',
path_to_center_param)
path_to_optimizer_center = cfg.MODEL.PRETRAIN_PATH.replace(
'model', 'optimizer_center')
print('Path to the checkpoint of optimizer_center:',
path_to_optimizer_center)
model.load_state_dict(torch.load(cfg.MODEL.PRETRAIN_PATH))
optimizer.load_state_dict(torch.load(path_to_optimizer))
center_criterion.load_state_dict(torch.load(path_to_center_param))
optimizer_center.load_state_dict(
torch.load(path_to_optimizer_center))
scheduler = WarmupMultiStepLR(optimizer, cfg.SOLVER.STEPS,
cfg.SOLVER.GAMMA,
cfg.SOLVER.WARMUP_FACTOR,
cfg.SOLVER.WARMUP_ITERS,
cfg.SOLVER.WARMUP_METHOD,
start_epoch)
elif cfg.MODEL.PRETRAIN_CHOICE == 'imagenet':
start_epoch = 0
scheduler = WarmupMultiStepLR(optimizer, cfg.SOLVER.STEPS,
cfg.SOLVER.GAMMA,
cfg.SOLVER.WARMUP_FACTOR,
cfg.SOLVER.WARMUP_ITERS,
cfg.SOLVER.WARMUP_METHOD)
else:
print(
'Only support pretrain_choice for imagenet and self, but got {}'
.format(cfg.MODEL.PRETRAIN_CHOICE))
do_train_with_center(
cfg,
model,
center_criterion,
train_loader,
val_loader,
optimizer,
optimizer_center,
scheduler, # modify for using self trained model
loss_func,
num_query,
start_epoch # add for using self trained model
)
else:
print(
"Unsupported value for cfg.MODEL.IF_WITH_CENTER {}, only support yes or no!\n"
.format(cfg.MODEL.IF_WITH_CENTER))
def main(args):
# transform
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
train_transform = transforms.Compose([
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
])
val_transform = transforms.Compose([
transforms.ToTensor(),
normalize
])
test_transform = transforms.Compose([
transforms.ToTensor(),
normalize
])
cap_transform = None
# data
train_loader = data_config(args.image_dir, args.anno_dir, args.batch_size, 'train', 100, train_transform, cap_transform=cap_transform)
test_loader = data_config(args.image_dir, args.anno_dir, 64, 'test', 100, test_transform)
unique_image = get_image_unique(args.image_dir, args.anno_dir, 64, 'test', 100, test_transform)
# loss
compute_loss = Loss(args)
nn.DataParallel(compute_loss).cuda()
# network
network, optimizer = network_config(args, 'train', compute_loss.parameters(), args.resume, args.model_path)
# lr_scheduler
scheduler = WarmupMultiStepLR(optimizer, (20, 25, 35), 0.1, 0.01, 10, 'linear')
ac_t2i_top1_best = 0.0
best_epoch = 0
for epoch in range(args.num_epoches - args.start_epoch):
network.train()
# train for one epoch
train_loss, train_time, image_precision, text_precision = train(args.start_epoch + epoch, train_loader, network, optimizer, compute_loss, args)
# evaluate on validation set
is_best = False
print('Train done for epoch-{}'.format(args.start_epoch + epoch))
logging.info('Epoch: [{}|{}], train_time: {:.3f}, train_loss: {:.3f}'.format(args.start_epoch + epoch, args.num_epoches, train_time, train_loss))
logging.info('image_precision: {:.3f}, text_precision: {:.3f}'.format(image_precision, text_precision))
scheduler.step()
for param in optimizer.param_groups:
print('lr:{}'.format(param['lr']))
if epoch >= 0:
ac_top1_i2t, ac_top5_i2t, ac_top10_i2t, ac_top1_t2i, ac_top5_t2i , ac_top10_t2i, test_time = test(test_loader, network, args, unique_image)
state = {'network': network.state_dict(), 'optimizer': optimizer.state_dict(), 'W': compute_loss.W, 'epoch': args.start_epoch + epoch}
if ac_top1_t2i > ac_t2i_top1_best:
best_epoch = epoch
ac_t2i_top1_best = ac_top1_t2i
save_checkpoint(state, epoch, args.checkpoint_dir, is_best)
logging.info('epoch:{}'.format(epoch))
logging.info('top1_t2i: {:.3f}, top5_t2i: {:.3f}, top10_t2i: {:.3f}, top1_i2t: {:.3f}, top5_i2t: {:.3f}, top10_i2t: {:.3f}'.format(
ac_top1_t2i, ac_top5_t2i, ac_top10_t2i, ac_top1_i2t, ac_top5_i2t, ac_top10_i2t))
logging.info('Best epoch:{}'.format(best_epoch))
logging.info('Train done')
logging.info(args.checkpoint_dir)
logging.info(args.log_dir)
class BaseModel(object):
def __init__(self, cfg):
self.cfg = cfg
self._init_models()
self._init_optimizers()
print('---------- Networks initialized -------------')
print_network(self.Content_Encoder)
print('-----------------------------------------------')
def _init_models(self):
# -----------------Content_Encoder-------------------
self.Content_Encoder = Baseline(self.cfg.DATASETS.NUM_CLASSES_S, 1, self.cfg.MODEL.PRETRAIN_PATH, 'bnneck',
'after', self.cfg.MODEL.NAME, 'imagenet')
# -----------------Criterion----------------- #
self.xent = CrossEntropyLabelSmooth(num_classes=self.cfg.DATASETS.NUM_CLASSES_S).cuda()
self.triplet = TripletLoss(0.3)
self.Smooth_L1_loss = torch.nn.SmoothL1Loss(reduction='mean').cuda()
# --------------------Cuda------------------- #
self.Content_Encoder = torch.nn.DataParallel(self.Content_Encoder).cuda()
def _init_optimizers(self):
self.Content_optimizer = make_optimizer(self.cfg, self.Content_Encoder)
self.Content_optimizer_fix = make_optimizer(self.cfg, self.Content_Encoder, fix=True)
self.scheduler = WarmupMultiStepLR(self.Content_optimizer, (30, 55), 0.1, 1.0 / 3,
500, "linear")
self.scheduler_fix = WarmupMultiStepLR(self.Content_optimizer_fix, (30, 55), 0.1, 1.0 / 3,
500, "linear")
self.schedulers = []
self.optimizers = []
def reset_model_status(self):
self.Content_Encoder.train()
def two_classifier(self, epoch, train_loader_s, train_loader_t, writer, logger, rand_src_1, rand_src_2,
print_freq=1):
self.reset_model_status()
self.epoch = epoch
self.scheduler.step(epoch)
self.scheduler_fix.step(epoch)
target_iter = iter(train_loader_t)
batch_time = AverageMeter()
data_time = AverageMeter()
end = time.time()
if (epoch < 80) or (110 <= epoch < 170):
mode = 'normal_c1_c2'
elif (80 <= epoch < 110) or (170 <= epoch < 210):
mode = 'reverse_c1_c2'
elif 210 <= epoch:
mode = 'fix_c1_c2'
for i, inputs in enumerate(train_loader_s):
data_time.update(time.time() - end)
try:
inputs_target = next(target_iter)
except:
target_iter = iter(train_loader_t)
inputs_target = next(target_iter)
img_s, pid_s, camid_s = self._parse_data(inputs)
img_t, pid_t, camid_t = self._parse_data(inputs_target)
content_code_s, content_feat_s = self.Content_Encoder(img_s)
pid_s_12 = np.asarray(pid_s.cpu())
camid_s = np.asarray(camid_s.cpu())
idx = []
for c_id in rand_src_1:
if len(np.where(c_id == camid_s)[0]) == 0:
continue
else:
idx.append(np.where(c_id == camid_s)[0])
if idx == [] or len(idx[0]) == 1:
idx = [np.asarray([a]) for a in range(self.cfg.SOLVER.IMS_PER_BATCH)]
idx = np.concatenate(idx)
pid_1 = torch.tensor(pid_s_12[idx]).cuda()
feat_1 = content_feat_s[idx]
idx = []
for c_id in rand_src_2:
if len(np.where(c_id == camid_s)[0]) == 0:
continue
else:
idx.append(np.where(c_id == camid_s)[0])
if idx == [] or len(idx[0]) == 1:
idx = [np.asarray([a]) for a in range(self.cfg.SOLVER.IMS_PER_BATCH)]
idx = np.concatenate(idx)
pid_2 = torch.tensor(pid_s_12[idx]).cuda()
feat_2 = content_feat_s[idx]
if mode == 'normal_c1_c2':
class_1 = self.Content_Encoder(feat_1, mode='c1')
class_2 = self.Content_Encoder(feat_2, mode='c2')
ID_loss_1 = self.xent(class_1, pid_1)
ID_loss_2 = self.xent(class_2, pid_2)
ID_tri_loss = self.triplet(content_feat_s, pid_s)
total_loss = ID_loss_1 + ID_loss_2 + ID_tri_loss[0]
self.Content_optimizer.zero_grad()
total_loss.backward()
self.Content_optimizer.step()
batch_time.update(time.time() - end)
end = time.time()
if (i + 1) % print_freq == 0:
logger.info('Epoch: [{}][{}/{}]\t'
'Time {:.3f} ({:.3f})\t'
'Data {:.3f} ({:.3f})\t'
'ID_loss: {:.3f} ID_loss_1: {:.3f} ID_loss_2: {:.3f} tri_loss: {:.3f} '
.format(epoch, i + 1, len(train_loader_s),
batch_time.val, batch_time.avg,
data_time.val, data_time.avg,
total_loss.item(), ID_loss_1.item(), ID_loss_2.item(), ID_tri_loss[0].item()
))
elif mode == 'reverse_c1_c2':
class_1 = self.Content_Encoder(feat_1, mode='c2')
class_2 = self.Content_Encoder(feat_2, mode='c1')
ID_loss_1 = self.xent(class_1, pid_1)
ID_loss_2 = self.xent(class_2, pid_2)
ID_tri_loss = self.triplet(content_feat_s, pid_s)
total_loss = ID_loss_1 + ID_loss_2 + ID_tri_loss[0]
self.Content_optimizer_fix.zero_grad()
total_loss.backward()
self.Content_optimizer_fix.step()
batch_time.update(time.time() - end)
end = time.time()
if (i + 1) % print_freq == 0:
logger.info('Epoch: [{}][{}/{}]\t'
'Time {:.3f} ({:.3f})\t'
'Data {:.3f} ({:.3f})\t'
'ID_loss: {:.3f} ID_loss_1: {:.3f} ID_loss_2: {:.3f} tri_loss: {:.3f}'
.format(epoch, i + 1, len(train_loader_s),
batch_time.val, batch_time.avg,
data_time.val, data_time.avg,
total_loss.item(), ID_loss_1.item(), ID_loss_2.item(), ID_tri_loss[0].item()
))
elif mode == 'fix_c1_c2':
class_1 = self.Content_Encoder(feat_1, mode='c2')
class_2 = self.Content_Encoder(feat_2, mode='c1')
ID_loss_1 = self.xent(class_1, pid_1)
ID_loss_2 = self.xent(class_2, pid_2)
content_code_t, content_feat_t = self.Content_Encoder(img_t)
tar_class_1 = self.Content_Encoder(content_feat_t, mode='c1')
tar_class_2 = self.Content_Encoder(content_feat_t, mode='c2')
tar_L1_loss = self.Smooth_L1_loss(tar_class_1, tar_class_2)
ID_tri_loss = self.triplet(content_feat_s, pid_s)
arg_c1 = torch.argmax(tar_class_1, dim=1)
arg_c2 = torch.argmax(tar_class_2, dim=1)
arg_idx = []
fake_id = []
for i_dx, data in enumerate(arg_c1):
if (data == arg_c2[i_dx]) and (((tar_class_1[i_dx][data] + tar_class_2[i_dx][arg_c2[i_dx]])/2) > 0.8):
arg_idx.append(i_dx)
fake_id.append(data)
if 210 <= epoch < 220:
if arg_idx != []:
ID_loss_fake = self.xent(content_code_t[arg_idx], torch.tensor(fake_id).cuda())
total_loss = ID_loss_1 + ID_loss_2 + 0.5 * tar_L1_loss + ID_tri_loss[0]
else:
ID_loss_fake = torch.tensor([0])
total_loss = ID_loss_1 + ID_loss_2 + 0.5 * tar_L1_loss + ID_tri_loss[0]
if 220 <= epoch:
if arg_idx != []:
ID_loss_fake = self.xent(content_code_t[arg_idx], torch.tensor(fake_id).cuda())
total_loss = ID_loss_1 + ID_loss_2 + 0.08 * ID_loss_fake + ID_tri_loss[0] + 0.5 * tar_L1_loss
else:
ID_loss_fake = torch.tensor([0])
total_loss = ID_loss_1 + ID_loss_2 + ID_tri_loss[0] + 0.5 * tar_L1_loss
self.Content_optimizer_fix.zero_grad()
total_loss.backward()
self.Content_optimizer_fix.step()
batch_time.update(time.time() - end)
end = time.time()
if (i + 1) % print_freq == 0:
logger.info('Epoch: [{}][{}/{}]\t'
'Time {:.3f} ({:.3f})\t'
'Data {:.3f} ({:.3f})\t'
'ID_loss: {:.3f} ID_loss_1: {:.3f} ID_loss_2: {:.3f} tar_L1_loss: {:.3f} tri_loss: {:.3f} ID_loss_fake: {:.6f}'
.format(epoch, i + 1, len(train_loader_s),
batch_time.val, batch_time.avg,
data_time.val, data_time.avg,
total_loss.item(), ID_loss_1.item(), ID_loss_2.item(), tar_L1_loss.item(),
ID_tri_loss[0].item(), ID_loss_fake.item()))
def _parse_data(self, inputs):
imgs, pids, camids = inputs
inputs = imgs.cuda()
targets = pids.cuda()
camids = camids.cuda()
return inputs, targets, camids
#CONFIG PARSER
config = get_args()
output_path = config.output_path
make_dirs(output_path)
logger = setup_logger('reid_baseline',output_path,if_train=True)
train_loader,train_gen_loader, val_loader, num_query, num_classes = make_dataloader(config)
model = Backbone(num_classes,config)
# if config.pretrain:
# model.load_param_finetune(config.m_pretrain_path)
loss_func, center_criterion = make_loss(config, num_classes=num_classes)
optimizer, optimizer_center = make_optimizer( model, center_criterion)
scheduler = WarmupMultiStepLR(optimizer, [40,70], 0.1,
0.01,
10, 'linear')
log_period = config.log_interval
checkpoint_period = config.save_model_interval
eval_period = config.test_interval
device = "cuda"
epochs = 80
logger = logging.getLogger("reid_baseline.train")
logger.info('start training')
if device:
if torch.cuda.device_count() > 1:
print('Using {} GPUs for training'.format(torch.cuda.device_count()))
def train(cfg):
# prepare dataset
# train_loader, val_loader, num_query, num_classes = make_data_loader(cfg)
train_loader, val_loader, num_query, num_classes = make_data_loader_train(
cfg)
# prepare model
if 'prw' in cfg.DATASETS.NAMES:
num_classes = 483
elif "market1501" in cfg.DATASETS.NAMES:
num_classes = 751
elif "duke" in cfg.DATASETS.NAMES:
num_classes = 702
elif "cuhk" in cfg.DATASETS.NAMES:
num_classes = 5532
model = build_model(cfg, num_classes)
if cfg.MODEL.IF_WITH_CENTER == 'no':
print('Train without center loss, the loss type is',
cfg.MODEL.METRIC_LOSS_TYPE)
optimizer = make_optimizer(cfg, model)
# scheduler = WarmupMultiStepLR(optimizer, cfg.SOLVER.STEPS, cfg.SOLVER.GAMMA, cfg.SOLVER.WARMUP_FACTOR,
# cfg.SOLVER.WARMUP_ITERS, cfg.SOLVER.WARMUP_METHOD)
loss_func = make_loss(cfg, num_classes) # modified by gu
# Add for using self trained model
if cfg.MODEL.PRETRAIN_CHOICE == 'self':
# start_epoch = eval(cfg.MODEL.PRETRAIN_PATH.split('/')[-1].split('.')[0].split('_')[-1])
start_epoch = 0
print('Start epoch:', start_epoch)
path_to_optimizer = cfg.MODEL.PRETRAIN_PATH.replace(
'model', 'optimizer')
print('Path to the checkpoint of optimizer:', path_to_optimizer)
pretrained_dic = torch.load(cfg.MODEL.PRETRAIN_PATH).state_dict()
model_dict = model.state_dict()
model_dict.update(pretrained_dic)
model.load_state_dict(model_dict)
if cfg.MODEL.WHOLE_MODEL_TRAIN == "no":
for name, value in model.named_parameters():
if "Query_Guided_Attention" not in name and "non_local" not in name and "classifier_attention" not in name:
value.requires_grad = False
optimizer = make_optimizer(cfg, model)
# else:
# cfg.SOLVER.BASE_LR = 0.0000035
# optimizer.load_state_dict(torch.load(path_to_optimizer))
# #####
# for state in optimizer.state.values():
# for k, v in state.items():
# if isinstance(v, torch.Tensor):
# state[k] = v.cuda()
# #####
scheduler = WarmupMultiStepLR(optimizer, cfg.SOLVER.STEPS,
cfg.SOLVER.GAMMA,
cfg.SOLVER.WARMUP_FACTOR,
cfg.SOLVER.WARMUP_ITERS,
cfg.SOLVER.WARMUP_METHOD)
elif cfg.MODEL.PRETRAIN_CHOICE == 'imagenet':
start_epoch = 0
scheduler = WarmupMultiStepLR(optimizer, cfg.SOLVER.STEPS,
cfg.SOLVER.GAMMA,
cfg.SOLVER.WARMUP_FACTOR,
cfg.SOLVER.WARMUP_ITERS,
cfg.SOLVER.WARMUP_METHOD)
else:
print(
'Only support pretrain_choice for imagenet and self, but got {}'
.format(cfg.MODEL.PRETRAIN_CHOICE))
arguments = {}
do_train(
cfg,
model,
train_loader,
val_loader,
optimizer,
scheduler, # modify for using self trained model
loss_func,
num_query,
start_epoch # add for using self trained model
)
elif cfg.MODEL.IF_WITH_CENTER == 'yes':
print('Train with center loss, the loss type is',
cfg.MODEL.METRIC_LOSS_TYPE)
loss_func, center_criterion = make_loss_with_center(
cfg, num_classes) # modified by gu
optimizer, optimizer_center = make_optimizer_with_center(
cfg, model, center_criterion)
# scheduler = WarmupMultiStepLR(optimizer, cfg.SOLVER.STEPS, cfg.SOLVER.GAMMA, cfg.SOLVER.WARMUP_FACTOR,
# cfg.SOLVER.WARMUP_ITERS, cfg.SOLVER.WARMUP_METHOD)
arguments = {}
# Add for using self trained model
if cfg.MODEL.PRETRAIN_CHOICE == 'self':
start_epoch = eval(
cfg.MODEL.PRETRAIN_PATH.split('/')[-1].split('.')[0].split('_')
[-1])
print('Start epoch:', start_epoch)
path_to_optimizer = cfg.MODEL.PRETRAIN_PATH.replace(
'model', 'optimizer')
print('Path to the checkpoint of optimizer:', path_to_optimizer)
path_to_center_param = cfg.MODEL.PRETRAIN_PATH.replace(
'model', 'center_param')
print('Path to the checkpoint of center_param:',
path_to_center_param)
path_to_optimizer_center = cfg.MODEL.PRETRAIN_PATH.replace(
'model', 'optimizer_center')
print('Path to the checkpoint of optimizer_center:',
path_to_optimizer_center)
model.load_state_dict(torch.load(cfg.MODEL.PRETRAIN_PATH))
optimizer.load_state_dict(torch.load(path_to_optimizer))
#####
for state in optimizer.state.values():
for k, v in state.items():
if isinstance(v, torch.Tensor):
state[k] = v.cuda()
#####
center_criterion.load_state_dict(torch.load(path_to_center_param))
optimizer_center.load_state_dict(
torch.load(path_to_optimizer_center))
scheduler = WarmupMultiStepLR(optimizer, cfg.SOLVER.STEPS,
cfg.SOLVER.GAMMA,
cfg.SOLVER.WARMUP_FACTOR,
cfg.SOLVER.WARMUP_ITERS,
cfg.SOLVER.WARMUP_METHOD,
start_epoch)
elif cfg.MODEL.PRETRAIN_CHOICE == 'imagenet':
start_epoch = 0
scheduler = WarmupMultiStepLR(optimizer, cfg.SOLVER.STEPS,
cfg.SOLVER.GAMMA,
cfg.SOLVER.WARMUP_FACTOR,
cfg.SOLVER.WARMUP_ITERS,
cfg.SOLVER.WARMUP_METHOD)
else:
print(
'Only support pretrain_choice for imagenet and self, but got {}'
.format(cfg.MODEL.PRETRAIN_CHOICE))
do_train_with_center(
cfg,
model,
center_criterion,
train_loader,
val_loader,
optimizer,
optimizer_center,
scheduler, # modify for using self trained model
loss_func,
num_query,
start_epoch # add for using self trained model
)
else:
print(
"Unsupported value for cfg.MODEL.IF_WITH_CENTER {}, only support yes or no!\n"
.format(cfg.MODEL.IF_WITH_CENTER))
def train(cfg):
# prepare dataset
train_loader, val_loader, num_query, num_classes, num_classes2, image_map_label2 = make_data_loader(
cfg)
#print('\n\n*** image_map_label2:')
# prepare model
model = build_model(cfg, num_classes, num_classes2)
#print(list(model.children()))
#print(model.state_dict().keys())
#exit(0)
#print('model.named_children(): \n\n', model.named_children())
'''
kk = 1
for name, child in model.base.named_children():
print(kk, name)
kk += 1
print(len(list(model.base.children())))
exit(0)
for i in range(len(list(model.base.children()))):
print(' +++', i+1)
print(list(model.base.children())[i])
exit(0)
'''
if len(cfg.MODEL.PRETRAIN_PATH2) > 5:
print('--- resume from ', cfg.MODEL.PRETRAIN_PATH2)
#model.load_param(cfg.MODEL.PRETRAIN_PATH)
#model.loiad_state_dict(torch.load(cfg.MODEL.PRETRAIN_PATH2, map_location=lambda storage, loc: storage))
if cfg.MODEL.ONCE_LOAD == 'yes':
print('\n---ONCE_LOAD...\n')
model.load_state_dict(
torch.load(cfg.MODEL.PRETRAIN_PATH2,
map_location=lambda storage, loc: storage))
#if cfg.MODEL.FREEZE_BASE == 'yes':
# functions.freeze_layer(model, 'base', False)
#functions.freeze_global_model(model, False)
else:
functions.load_state_dict_distill(model, cfg.MODEL.PRETRAIN_PATH2,
cfg.MODEL.ONLY_BASE,
cfg.MODEL.WITHOUT_FC)
print('**** Successfully load ', cfg.MODEL.PRETRAIN_PATH2)
if cfg.MODEL.FREEZE_BASE:
#functions.freeze_layer(model, 'base', False)
functions.freeze_global_model(model, False)
if cfg.MODEL.IF_WITH_CENTER == 'no':
print('Train without center loss, the loss type is',
cfg.MODEL.METRIC_LOSS_TYPE)
if cfg.SOLVER.MY_OPTIMIZER == "yes":
print('---* my optimizer:', cfg.SOLVER.MY_OPTIMIZER_NAME)
other_params = [
p for n, p in model.named_parameters()
if not n.startswith('base')
]
optimizer = optim.SGD([{
'params': model.base.parameters(),
'lr': cfg.SOLVER.LR / 10
}, {
'params': other_params,
'lr': cfg.SOLVER.LR
}],
momentum=0.9,
weight_decay=5e-4,
nesterov=True)
else:
print('---* not my optimizer')
optimizer = make_optimizer(cfg, model)
# scheduler = WarmupMultiStepLR(optimizer, cfg.SOLVER.STEPS, cfg.SOLVER.GAMMA, cfg.SOLVER.WARMUP_FACTOR,
# cfg.SOLVER.WARMUP_ITERS, cfg.SOLVER.WARMUP_METHOD)
#_C.SOLVER.MY_SCHEDULER = "no"
#_C.SOLVER.MY_WARMUP = "no"
loss_func = make_loss(cfg, num_classes) # modified by gu
# Add for using self trained model
if cfg.MODEL.PRETRAIN_CHOICE == 'self':
start_epoch = eval(
cfg.MODEL.PRETRAIN_PATH.split('/')[-1].split('.')[0].split('_')
[-1])
print('Start epoch:', start_epoch)
path_to_optimizer = cfg.MODEL.PRETRAIN_PATH.replace(
'model', 'optimizer')
print('Path to the checkpoint of optimizer:', path_to_optimizer)
model.load_state_dict(torch.load(cfg.MODEL.PRETRAIN_PATH))
optimizer.load_state_dict(torch.load(path_to_optimizer))
scheduler = WarmupMultiStepLR(optimizer, cfg.SOLVER.STEPS,
cfg.SOLVER.GAMMA,
cfg.SOLVER.WARMUP_FACTOR,
cfg.SOLVER.WARMUP_ITERS,
cfg.SOLVER.WARMUP_METHOD,
start_epoch)
elif cfg.MODEL.PRETRAIN_CHOICE == 'imagenet':
start_epoch = 0
if cfg.SOLVER.MY_SCHEDULER == "yes":
print('cfg.SOLVER.MY_SCHEDULER_STEP:',
cfg.SOLVER.MY_SCHEDULER_STEP)
print('---* my scheduler: ', cfg.SOLVER.MY_SCHEDULER_NAME)
if cfg.SOLVER.MY_SCHEDULER_NAME == 'SL':
scheduler = lr_scheduler.StepLR(
optimizer,
step_size=cfg.SOLVER.MY_SCHEDULER_STEP[0],
gamma=0.1)
elif cfg.SOLVER.MY_SCHEDULER_NAME == 'MSL':
scheduler = lr_scheduler.MultiStepLR(
optimizer, cfg.SOLVER.MY_SCHEDULER_STEP, gamma=0.1)
else:
print(cfg.SOLVER.MY_SCHEDULER_NAME, ' not found!')
eixt(0)
else:
print('---* not my scheduler')
scheduler = WarmupMultiStepLR(optimizer, cfg.SOLVER.STEPS,
cfg.SOLVER.GAMMA,
cfg.SOLVER.WARMUP_FACTOR,
cfg.SOLVER.WARMUP_ITERS,
cfg.SOLVER.WARMUP_METHOD)
else:
print(
'Only support pretrain_choice for imagenet and self, but got {}'
.format(cfg.MODEL.PRETRAIN_CHOICE))
arguments = {}
print('************ do_train')
do_train(
cfg,
model,
train_loader,
val_loader,
optimizer,
scheduler, # modify for using self trained model
loss_func,
num_query,
start_epoch, # add for using self trained model
image_map_label2,
num_classes2)
# elif cfg.MODEL.IF_WITH_CENTER == 'yes':
# print('Train with center loss, the loss type is', cfg.MODEL.METRIC_LOSS_TYPE)
# loss_func, center_criterion = make_loss_with_center(cfg, num_classes) # modified by gu
# optimizer, optimizer_center = make_optimizer_with_center(cfg, model, center_criterion)
# # scheduler = WarmupMultiStepLR(optimizer, cfg.SOLVER.STEPS, cfg.SOLVER.GAMMA, cfg.SOLVER.WARMUP_FACTOR,
# # cfg.SOLVER.WARMUP_ITERS, cfg.SOLVER.WARMUP_METHOD)
#
# arguments = {}
#
# # Add for using self trained model
# if cfg.MODEL.PRETRAIN_CHOICE == 'self':
# start_epoch = eval(cfg.MODEL.PRETRAIN_PATH.split('/')[-1].split('.')[0].split('_')[-1])
# print('Start epoch:', start_epoch)
# path_to_optimizer = cfg.MODEL.PRETRAIN_PATH.replace('model', 'optimizer')
# print('Path to the checkpoint of optimizer:', path_to_optimizer)
# path_to_optimizer_center = cfg.MODEL.PRETRAIN_PATH.replace('model', 'optimizer_center')
# print('Path to the checkpoint of optimizer_center:', path_to_optimizer_center)
# model.load_state_dict(torch.load(cfg.MODEL.PRETRAIN_PATH))
# optimizer.load_state_dict(torch.load(path_to_optimizer))
# optimizer_center.load_state_dict(torch.load(path_to_optimizer_center))
# scheduler = WarmupMultiStepLR(optimizer, cfg.SOLVER.STEPS, cfg.SOLVER.GAMMA, cfg.SOLVER.WARMUP_FACTOR,
# cfg.SOLVER.WARMUP_ITERS, cfg.SOLVER.WARMUP_METHOD, start_epoch)
# elif cfg.MODEL.PRETRAIN_CHOICE == 'imagenet':
# start_epoch = 0
# scheduler = WarmupMultiStepLR(optimizer, cfg.SOLVER.STEPS, cfg.SOLVER.GAMMA, cfg.SOLVER.WARMUP_FACTOR,
# cfg.SOLVER.WARMUP_ITERS, cfg.SOLVER.WARMUP_METHOD)
# else:
# print('Only support pretrain_choice for imagenet and self, but got {}'.format(cfg.MODEL.PRETRAIN_CHOICE))
#
# do_train_with_center(
# cfg,
# model,
# center_criterion,
# train_loader,
# val_loader,
# optimizer,
# optimizer_center,
# scheduler, # modify for using self trained model
# loss_func,
# num_query,
# start_epoch # add for using self trained model
# )
else:
print(
"Unsupported value for cfg.MODEL.IF_WITH_CENTER {}, only support yes or no!\n"
.format(cfg.MODEL.IF_WITH_CENTER))
def train(cfg):
# prepare dataset
train_loader, val_loader, num_query, num_classes = make_data_loader(cfg)
# prepare model
model = build_model(cfg, num_classes)
if cfg.MODEL.IF_WITH_CENTER == 'no':
print('Train without center loss, the loss type is',
cfg.MODEL.METRIC_LOSS_TYPE)
optimizer = make_optimizer(cfg, model)
# scheduler = WarmupMultiStepLR(optimizer, cfg.SOLVER.STEPS, cfg.SOLVER.GAMMA, cfg.SOLVER.WARMUP_FACTOR,
# cfg.SOLVER.WARMUP_ITERS, cfg.SOLVER.WARMUP_METHOD)
loss_func = make_loss(cfg, num_classes) # modified by gu
# Add for using self trained model
if cfg.MODEL.PRETRAIN_CHOICE == 'self':
start_epoch = eval(
cfg.MODEL.PRETRAIN_PATH.split('/')[-1].split('.')[0].split('_')
[-1])
print('Start epoch:', start_epoch)
path_to_optimizer = cfg.MODEL.PRETRAIN_PATH.replace(
'model', 'optimizer')
print('Path to the checkpoint of optimizer:', path_to_optimizer)
model.load_state_dict(torch.load(cfg.MODEL.PRETRAIN_PATH))
optimizer.load_state_dict(torch.load(path_to_optimizer))
scheduler = WarmupMultiStepLR(optimizer, cfg.SOLVER.STEPS,
cfg.SOLVER.GAMMA,
cfg.SOLVER.WARMUP_FACTOR,
cfg.SOLVER.WARMUP_ITERS,
cfg.SOLVER.WARMUP_METHOD,
start_epoch)
elif cfg.MODEL.PRETRAIN_CHOICE == 'imagenet':
start_epoch = 0
scheduler = WarmupMultiStepLR(optimizer, cfg.SOLVER.STEPS,
cfg.SOLVER.GAMMA,
cfg.SOLVER.WARMUP_FACTOR,
cfg.SOLVER.WARMUP_ITERS,
cfg.SOLVER.WARMUP_METHOD)
else:
print(
'Only support pretrain_choice for imagenet and self, but got {}'
.format(cfg.MODEL.PRETRAIN_CHOICE))
arguments = {}
do_train(
cfg,
model,
train_loader,
val_loader,
optimizer,
scheduler, # modify for using self trained model
loss_func,
num_query,
start_epoch # add for using self trained model
)
elif cfg.MODEL.IF_WITH_CENTER == 'yes':
print('Train with center loss, the loss type is',
cfg.MODEL.METRIC_LOSS_TYPE)
loss_func, center_criterion = make_loss_with_center(
cfg, num_classes) # modified by gu
optimizer, optimizer_center = make_optimizer_with_center(
cfg, model, center_criterion)
# scheduler = WarmupMultiStepLR(optimizer, cfg.SOLVER.STEPS, cfg.SOLVER.GAMMA, cfg.SOLVER.WARMUP_FACTOR,
# cfg.SOLVER.WARMUP_ITERS, cfg.SOLVER.WARMUP_METHOD)
arguments = {}
# Add for using self trained model
if cfg.MODEL.PRETRAIN_CHOICE == 'self':
start_epoch = eval(
cfg.MODEL.PRETRAIN_PATH.split('/')[-1].split('.')[0].split('_')
[-1])
print('Start epoch:', start_epoch)
path_to_optimizer = cfg.MODEL.PRETRAIN_PATH.replace(
'model', 'optimizer')
print('Path to the checkpoint of optimizer:', path_to_optimizer)
path_to_center_param = cfg.MODEL.PRETRAIN_PATH.replace(
'model', 'center_param')
print('Path to the checkpoint of center_param:',
path_to_center_param)
path_to_optimizer_center = cfg.MODEL.PRETRAIN_PATH.replace(
'model', 'optimizer_center')
print('Path to the checkpoint of optimizer_center:',
path_to_optimizer_center)
model.load_state_dict(torch.load(cfg.MODEL.PRETRAIN_PATH))
optimizer.load_state_dict(torch.load(path_to_optimizer))
center_criterion.load_state_dict(torch.load(path_to_center_param))
optimizer_center.load_state_dict(
torch.load(path_to_optimizer_center))
scheduler = WarmupMultiStepLR(optimizer, cfg.SOLVER.STEPS,
cfg.SOLVER.GAMMA,
cfg.SOLVER.WARMUP_FACTOR,
cfg.SOLVER.WARMUP_ITERS,
cfg.SOLVER.WARMUP_METHOD,
start_epoch)
elif cfg.MODEL.PRETRAIN_CHOICE == 'imagenet':
start_epoch = 0
scheduler = WarmupMultiStepLR(optimizer, cfg.SOLVER.STEPS,
cfg.SOLVER.GAMMA,
cfg.SOLVER.WARMUP_FACTOR,
cfg.SOLVER.WARMUP_ITERS,
cfg.SOLVER.WARMUP_METHOD)
else:
print(
'Only support pretrain_choice for imagenet and self, but got {}'
.format(cfg.MODEL.PRETRAIN_CHOICE))
do_train_with_center(
cfg,
model,
center_criterion,
train_loader,
val_loader,
optimizer,
optimizer_center,
scheduler, # modify for using self trained model
loss_func,
num_query,
start_epoch # add for using self trained model
)
else:
print(
"Unsupported value for cfg.MODEL.IF_WITH_CENTER {}, only support yes or no!\n"
.format(cfg.MODEL.IF_WITH_CENTER))
def train(cfg):
# prepare dataset
train_loader, val_loader, num_query, num_classes = make_data_loader(cfg)
# prepare model
model = build_model(cfg, num_classes)
if cfg.MODEL.IF_WITH_CENTER == "no":
print("Train without center loss, the loss type is", cfg.MODEL.METRIC_LOSS_TYPE)
optimizer = make_optimizer(cfg, model)
# scheduler = WarmupMultiStepLR(optimizer, cfg.SOLVER.STEPS, cfg.SOLVER.GAMMA, cfg.SOLVER.WARMUP_FACTOR,
# cfg.SOLVER.WARMUP_ITERS, cfg.SOLVER.WARMUP_METHOD)
loss_func = make_loss(cfg, num_classes) # modified by gu
# Add for using self trained model
if cfg.MODEL.PRETRAIN_CHOICE == "self":
start_epoch = eval(
cfg.MODEL.PRETRAIN_PATH.split("/")[-1].split(".")[0].split("_")[-1]
)
print("Start epoch:", start_epoch)
path_to_optimizer = cfg.MODEL.PRETRAIN_PATH.replace("model", "optimizer")
print("Path to the checkpoint of optimizer:", path_to_optimizer)
model.load_state_dict(torch.load(cfg.MODEL.PRETRAIN_PATH))
optimizer.load_state_dict(torch.load(path_to_optimizer))
scheduler = WarmupMultiStepLR(
optimizer,
cfg.SOLVER.STEPS,
cfg.SOLVER.GAMMA,
cfg.SOLVER.WARMUP_FACTOR,
cfg.SOLVER.WARMUP_ITERS,
cfg.SOLVER.WARMUP_METHOD,
cfg.SOLVER.MODE,
cfg.SOLVER.MAX_EPOCHS,
start_epoch
)
else:
start_epoch = 0
scheduler = WarmupMultiStepLR(
optimizer,
cfg.SOLVER.STEPS,
cfg.SOLVER.GAMMA,
cfg.SOLVER.WARMUP_FACTOR,
cfg.SOLVER.WARMUP_ITERS,
cfg.SOLVER.WARMUP_METHOD,
cfg.SOLVER.MODE,
cfg.SOLVER.MAX_EPOCHS
)
do_train(
cfg,
model,
train_loader,
val_loader,
optimizer,
scheduler, # modify for using self trained model
loss_func,
num_query,
start_epoch, # add for using self trained model
)
elif cfg.MODEL.IF_WITH_CENTER == "yes":
print("Train with center loss, the loss type is", cfg.MODEL.METRIC_LOSS_TYPE)
loss_func, center_criterion = make_loss_with_center(
cfg, num_classes
) # modified by gu
optimizer, optimizer_center = make_optimizer_with_center(
cfg, model, center_criterion
)
# scheduler = WarmupMultiStepLR(optimizer, cfg.SOLVER.STEPS, cfg.SOLVER.GAMMA, cfg.SOLVER.WARMUP_FACTOR,
# cfg.SOLVER.WARMUP_ITERS, cfg.SOLVER.WARMUP_METHOD)
# Add for using self trained model
if cfg.MODEL.PRETRAIN_CHOICE == "self":
start_epoch = eval(
cfg.MODEL.PRETRAIN_PATH.split("/")[-1].split(".")[0].split("_")[-1]
)
print("Start epoch:", start_epoch)
path_to_optimizer = cfg.MODEL.PRETRAIN_PATH.replace("model", "optimizer")
print("Path to the checkpoint of optimizer:", path_to_optimizer)
path_to_center_param = cfg.MODEL.PRETRAIN_PATH.replace(
"model", "center_param"
)
print("Path to the checkpoint of center_param:", path_to_center_param)
path_to_optimizer_center = cfg.MODEL.PRETRAIN_PATH.replace(
"model", "optimizer_center"
)
print(
"Path to the checkpoint of optimizer_center:", path_to_optimizer_center
)
model.load_state_dict(torch.load(cfg.MODEL.PRETRAIN_PATH))
optimizer.load_state_dict(torch.load(path_to_optimizer))
center_criterion.load_state_dict(torch.load(path_to_center_param))
optimizer_center.load_state_dict(torch.load(path_to_optimizer_center))
scheduler = WarmupMultiStepLR(
optimizer,
cfg.SOLVER.STEPS,
cfg.SOLVER.GAMMA,
cfg.SOLVER.WARMUP_FACTOR,
cfg.SOLVER.WARMUP_ITERS,
cfg.SOLVER.WARMUP_METHOD,
cfg.SOLVER.MODE,
cfg.SOLVER.MAX_EPOCHS,
start_epoch,
)
else:
start_epoch = 0
scheduler = WarmupMultiStepLR(
optimizer,
cfg.SOLVER.STEPS,
cfg.SOLVER.GAMMA,
cfg.SOLVER.WARMUP_FACTOR,
cfg.SOLVER.WARMUP_ITERS,
cfg.SOLVER.WARMUP_METHOD,
cfg.SOLVER.MODE,
cfg.SOLVER.MAX_EPOCHS,
)
do_train_with_center(
cfg,
model,
center_criterion,
train_loader,
val_loader,
optimizer,
optimizer_center,
scheduler, # modify for using self trained model
loss_func,
num_query,
start_epoch, # add for using self trained model
)
else:
print(
"Unsupported value for cfg.MODEL.IF_WITH_CENTER {}, only support yes or no!\n".format(
cfg.MODEL.IF_WITH_CENTER
)
)
def train(cfg):
# prepare dataset
train_loader, val_loader, num_query, num_classes = make_data_loader(cfg)
# prepare model
model = build_model(cfg, num_classes)
print('Train with the loss type is', cfg.MODEL.METRIC_LOSS_TYPE)
optimizer = make_optimizer(cfg, model)
# scheduler = WarmupMultiStepLR(optimizer, cfg.SOLVER.STEPS, cfg.SOLVER.GAMMA, cfg.SOLVER.WARMUP_FACTOR,
# cfg.SOLVER.WARMUP_ITERS, cfg.SOLVER.WARMUP_METHOD)
loss_func = make_loss(cfg, num_classes) # modified by gu
# Add for using self trained model
if cfg.MODEL.PRETRAIN_CHOICE == 'self':
# start_epoch = eval(cfg.MODEL.PRETRAIN_PATH.split('/')[-1].split('.')[0].split('_')[-1])
# print('Start epoch:', start_epoch)
# path_to_optimizer = cfg.MODEL.PRETRAIN_PATH.replace('model', 'optimizer')
# print('Path to the checkpoint of optimizer:', path_to_optimizer)
# model.load_state_dict(torch.load(cfg.MODEL.PRETRAIN_PATH))
# optimizer.load_state_dict(torch.load(path_to_optimizer))
# scheduler = WarmupMultiStepLR(optimizer, cfg.SOLVER.STEPS, cfg.SOLVER.GAMMA, cfg.SOLVER.WARMUP_FACTOR,
# cfg.SOLVER.WARMUP_ITERS, cfg.SOLVER.WARMUP_METHOD, start_epoch)
start_epoch = 120
print('Start epoch:', start_epoch)
model.load_state_dict(
torch.load('work_space_tri/se_resnet101_ibn_a_model_120.pth'))
optimizer.load_state_dict(
torch.load('work_space_tri/se_resnet101_ibn_a_optimizer_120.pth'))
scheduler = WarmupMultiStepLR(optimizer, cfg.SOLVER.STEPS,
cfg.SOLVER.GAMMA,
cfg.SOLVER.WARMUP_FACTOR,
cfg.SOLVER.WARMUP_ITERS,
cfg.SOLVER.WARMUP_METHOD, start_epoch)
elif cfg.MODEL.PRETRAIN_CHOICE == 'imagenet':
start_epoch = 0
scheduler = WarmupMultiStepLR(optimizer, cfg.SOLVER.STEPS,
cfg.SOLVER.GAMMA,
cfg.SOLVER.WARMUP_FACTOR,
cfg.SOLVER.WARMUP_ITERS,
cfg.SOLVER.WARMUP_METHOD)
else:
print('Only support pretrain_choice for imagenet and self, but got {}'.
format(cfg.MODEL.PRETRAIN_CHOICE))
arguments = {}
do_train(
cfg,
model,
train_loader,
val_loader,
optimizer,
scheduler, # modify for using self trained model
loss_func,
num_query,
start_epoch # add for using self trained model
)
os.environ['CUDA_VISIBLE_DEVICES'] = cfg.MODEL.DEVICE_ID
train_loader, val_loader, num_query, num_classes = make_dataloader(cfg)
if cfg.MODEL.PRETRAIN_CHOICE == 'finetune':
model = make_model(cfg, num_class=num_classes)
model.load_param_finetune(cfg.MODEL.PRETRAIN_PATH)
print('Loading pretrained model for finetuning......')
else:
model = make_model(cfg, num_class=num_classes)
loss_func, center_criterion = make_loss(cfg, num_classes=num_classes)
optimizer, optimizer_center = make_optimizer(cfg, model, center_criterion)
scheduler = WarmupMultiStepLR(optimizer, cfg.SOLVER.STEPS,
cfg.SOLVER.GAMMA, cfg.SOLVER.WARMUP_FACTOR,
cfg.SOLVER.WARMUP_EPOCHS,
cfg.SOLVER.WARMUP_METHOD)
do_train(
cfg,
model,
center_criterion,
train_loader,
val_loader,
optimizer,
optimizer_center,
scheduler, # modify for using self trained model
loss_func,
num_query)
def main():
torch.backends.cudnn.deterministic = True
cudnn.benchmark = True
#parser = argparse.ArgumentParser(description="ReID Baseline Training")
#parser.add_argument(
#"--config_file", default="", help="path to config file", type=str)
#parser.add_argument("opts", help="Modify config options using the command-line", default=None, nargs=argparse.REMAINDER)
#args = parser.parse_args()
config_file = 'configs/baseline_veri_r101_a.yml'
if config_file != "":
cfg.merge_from_file(config_file)
#cfg.merge_from_list(args.opts)
cfg.freeze()
output_dir = cfg.OUTPUT_DIR
if output_dir and not os.path.exists(output_dir):
os.makedirs(output_dir)
logger = setup_logger("reid_baseline", output_dir, if_train=True)
logger.info("Saving model in the path :{}".format(cfg.OUTPUT_DIR))
logger.info(config_file)
if config_file != "":
logger.info("Loaded configuration file {}".format(config_file))
with open(config_file, 'r') as cf:
config_str = "\n" + cf.read()
logger.info(config_str)
logger.info("Running with config:\n{}".format(cfg))
os.environ['CUDA_VISIBLE_DEVICES'] = cfg.MODEL.DEVICE_ID
path = 'D:/Python_SMU/Veri/verigms/gms/'
pkl = {}
entries = os.listdir(path)
for name in entries:
f = open((path + name), 'rb')
if name == 'featureMatrix.pkl':
s = name[0:13]
else:
s = name[0:3]
pkl[s] = pickle.load(f)
f.close
with open('cids.pkl', 'rb') as handle:
b = pickle.load(handle)
with open('index.pkl', 'rb') as handle:
c = pickle.load(handle)
train_transforms, val_transforms, dataset, train_set, val_set = make_dataset(
cfg, pkl_file='index.pkl')
num_workers = cfg.DATALOADER.NUM_WORKERS
num_classes = dataset.num_train_pids
#pkl_f = 'index.pkl'
pid = 0
pidx = {}
for img_path, pid, _, _ in dataset.train:
path = img_path.split('\\')[-1]
folder = path[1:4]
pidx[folder] = pid
pid += 1
if 'triplet' in cfg.DATALOADER.SAMPLER:
train_loader = DataLoader(train_set,
batch_size=cfg.SOLVER.IMS_PER_BATCH,
sampler=RandomIdentitySampler(
dataset.train, cfg.SOLVER.IMS_PER_BATCH,
cfg.DATALOADER.NUM_INSTANCE),
num_workers=num_workers,
pin_memory=True,
collate_fn=train_collate_fn)
elif cfg.DATALOADER.SAMPLER == 'softmax':
print('using softmax sampler')
train_loader = DataLoader(train_set,
batch_size=cfg.SOLVER.IMS_PER_BATCH,
shuffle=True,
num_workers=num_workers,
pin_memory=True,
collate_fn=train_collate_fn)
else:
print('unsupported sampler! expected softmax or triplet but got {}'.
format(cfg.SAMPLER))
print("train loader loaded successfully")
val_loader = DataLoader(val_set,
batch_size=cfg.TEST.IMS_PER_BATCH,
shuffle=False,
num_workers=num_workers,
pin_memory=True,
collate_fn=train_collate_fn)
print("val loader loaded successfully")
if cfg.MODEL.PRETRAIN_CHOICE == 'finetune':
model = make_model(cfg, num_class=576)
model.load_param_finetune(cfg.MODEL.PRETRAIN_PATH)
print('Loading pretrained model for finetuning......')
else:
model = make_model(cfg, num_class=num_classes)
loss_func, center_criterion = make_loss(cfg, num_classes=num_classes)
optimizer, optimizer_center = make_optimizer(cfg, model, center_criterion)
scheduler = WarmupMultiStepLR(optimizer, cfg.SOLVER.STEPS,
cfg.SOLVER.GAMMA, cfg.SOLVER.WARMUP_FACTOR,
cfg.SOLVER.WARMUP_EPOCHS,
cfg.SOLVER.WARMUP_METHOD)
print("model,optimizer, loss, scheduler loaded successfully")
height, width = cfg.INPUT.SIZE_TRAIN
log_period = cfg.SOLVER.LOG_PERIOD
checkpoint_period = cfg.SOLVER.CHECKPOINT_PERIOD
eval_period = cfg.SOLVER.EVAL_PERIOD
device = "cuda"
epochs = cfg.SOLVER.MAX_EPOCHS
logger = logging.getLogger("reid_baseline.train")
logger.info('start training')
if device:
if torch.cuda.device_count() > 1:
print('Using {} GPUs for training'.format(
torch.cuda.device_count()))
model = nn.DataParallel(model)
model.to(device)
loss_meter = AverageMeter()
acc_meter = AverageMeter()
evaluator = R1_mAP_eval(len(dataset.query),
max_rank=50,
feat_norm=cfg.TEST.FEAT_NORM)
model.base._freeze_stages()
logger.info('Freezing the stages number:{}'.format(cfg.MODEL.FROZEN))
data_index = search(pkl)
print("Ready for training")
for epoch in range(1, epochs + 1):
start_time = time.time()
loss_meter.reset()
acc_meter.reset()
evaluator.reset()
scheduler.step()
model.train()
for n_iter, (img, label, index, pid, cid) in enumerate(train_loader):
optimizer.zero_grad()
optimizer_center.zero_grad()
#img = img.to(device)
#target = vid.to(device)
trainX, trainY = torch.zeros(
(train_loader.batch_size * 3, 3, height, width),
dtype=torch.float32), torch.zeros(
(train_loader.batch_size * 3), dtype=torch.int64)
for i in range(train_loader.batch_size):
labelx = label[i]
indexx = index[i]
cidx = pid[i]
if indexx > len(pkl[labelx]) - 1:
indexx = len(pkl[labelx]) - 1
a = pkl[labelx][indexx]
minpos = np.argmin(ma.masked_where(a == 0, a))
pos_dic = train_set[data_index[cidx][1] + minpos]
#print(pos_dic[1])
neg_label = int(labelx)
while True:
neg_label = random.choice(range(1, 770))
if neg_label is not int(labelx) and os.path.isdir(
os.path.join('D:/datasets/veri-split/train',
strint(neg_label))) is True:
break
negative_label = strint(neg_label)
neg_cid = pidx[negative_label]
neg_index = random.choice(range(0, len(pkl[negative_label])))
neg_dic = train_set[data_index[neg_cid][1] + neg_index]
trainX[i] = img[i]
trainX[i + train_loader.batch_size] = pos_dic[0]
trainX[i + (train_loader.batch_size * 2)] = neg_dic[0]
trainY[i] = cidx
trainY[i + train_loader.batch_size] = pos_dic[3]
trainY[i + (train_loader.batch_size * 2)] = neg_dic[3]
#print(trainY)
trainX = trainX.cuda()
trainY = trainY.cuda()
score, feat = model(trainX, trainY)
loss = loss_func(score, feat, trainY)
loss.backward()
optimizer.step()
if 'center' in cfg.MODEL.METRIC_LOSS_TYPE:
for param in center_criterion.parameters():
param.grad.data *= (1. / cfg.SOLVER.CENTER_LOSS_WEIGHT)
optimizer_center.step()
acc = (score.max(1)[1] == trainY).float().mean()
loss_meter.update(loss.item(), img.shape[0])
acc_meter.update(acc, 1)
if (n_iter + 1) % log_period == 0:
logger.info(
"Epoch[{}] Iteration[{}/{}] Loss: {:.3f}, Acc: {:.3f}, Base Lr: {:.2e}"
.format(epoch, (n_iter + 1), len(train_loader),
loss_meter.avg, acc_meter.avg,
scheduler.get_lr()[0]))
end_time = time.time()
time_per_batch = (end_time - start_time) / (n_iter + 1)
logger.info(
"Epoch {} done. Time per batch: {:.3f}[s] Speed: {:.1f}[samples/s]"
.format(epoch, time_per_batch,
train_loader.batch_size / time_per_batch))
if epoch % checkpoint_period == 0:
torch.save(
model.state_dict(),
os.path.join(cfg.OUTPUT_DIR,
cfg.MODEL.NAME + '_{}.pth'.format(epoch)))
if epoch % eval_period == 0:
model.eval()
for n_iter, (img, vid, camid, _, _) in enumerate(val_loader):
with torch.no_grad():
img = img.to(device)
feat = model(img)
evaluator.update((feat, vid, camid))
cmc, mAP, _, _, _, _, _ = evaluator.compute()
logger.info("Validation Results - Epoch: {}".format(epoch))
logger.info("mAP: {:.1%}".format(mAP))
for r in [1, 5, 10]:
logger.info("CMC curve, Rank-{:<3}:{:.1%}".format(
r, cmc[r - 1]))
