def main():
# 1. Get input arguments
args = get_args()
# 2. Create config instance from args above
cfg = get_default_config()
cfg.use_gpu = torch.cuda.is_available()
if args.config_file:
cfg.merge_from_file(args.config_file)
reset_config(cfg, args)
cfg.merge_from_list(args.opts)
set_random_seed(cfg.train.seed)
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))
print('Collecting env info ...')
print('** System info **\n{}\n'.format(collect_env_info()))
if cfg.use_gpu:
torch.backends.cudnn.benchmark = True
# 3. Create DataManager Instance
datamanager = build_datamanager(cfg)
print('Building model: {}'.format(cfg.model.name))
model = torchreid.models.build_model(
name=cfg.model.name,
num_classes=datamanager.num_train_pids,
loss=cfg.loss.name,
pretrained=cfg.model.pretrained,
use_gpu=cfg.use_gpu)
num_params, flops = compute_model_complexity(
model, (1, 3, cfg.data.height, cfg.data.width))
print('Model complexity: params={:,} flops={:,}'.format(num_params, flops))
if cfg.model.load_weights and check_isfile(cfg.model.load_weights):
load_pretrained_weights(model, cfg.model.load_weights)
if cfg.use_gpu:
model = nn.DataParallel(model).cuda()
optimizer = torchreid.optim.build_optimizer(model, **optimizer_kwargs(cfg))
scheduler = torchreid.optim.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,
scheduler=scheduler)
print('Building {}-engine for {}-reid'.format(cfg.loss.name,
cfg.data.type))
# Build engine and run
engine = build_engine(cfg, datamanager, model, optimizer, scheduler)
engine.run(**engine_run_kwargs(cfg))
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() and not args.use_cpu
log_name = 'test.log' if args.evaluate else 'train.log'
log_name += time.strftime('-%Y-%m-%d-%H-%M-%S')
sys.stdout = Logger(osp.join(args.save_dir, log_name))
print('** Arguments **')
arg_keys = list(args.__dict__.keys())
arg_keys.sort()
for key in arg_keys:
print('{}: {}'.format(key, args.__dict__[key]))
print('\n')
print('Collecting env info ...')
print('** System info **\n{}\n'.format(collect_env_info()))
if use_gpu:
torch.backends.cudnn.benchmark = True
else:
warnings.warn(
'Currently using CPU, however, GPU is highly recommended')
datamanager = build_datamanager(args)
print('Building model: {}'.format(args.arch))
model = torchreid.models.build_model(
name=args.arch,
num_classes=datamanager.num_train_pids,
loss=args.loss.lower(),
pretrained=(not args.no_pretrained),
use_gpu=use_gpu)
num_params, flops = compute_model_complexity(
model, (1, 3, args.height, args.width))
print('Model complexity: params={:,} flops={:,}'.format(num_params, flops))
if args.load_weights and check_isfile(args.load_weights):
load_pretrained_weights(model, args.load_weights)
if use_gpu:
model = nn.DataParallel(model).cuda()
optimizer = torchreid.optim.build_optimizer(model,
**optimizer_kwargs(args))
scheduler = torchreid.optim.build_lr_scheduler(optimizer,
**lr_scheduler_kwargs(args))
if args.resume and check_isfile(args.resume):
args.start_epoch = resume_from_checkpoint(args.resume,
model,
optimizer=optimizer)
print('Building {}-engine for {}-reid'.format(args.loss, args.app))
engine = build_engine(args, datamanager, model, optimizer, scheduler)
engine.run(**engine_run_kwargs(args))
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('-s', '--sources', type=str, nargs='+', help='source datasets (delimited by space)')
parser.add_argument('-t', '--targets', type=str, nargs='+', help='target datasets (delimited by space)')
parser.add_argument('--root', type=str, default='', help='path to data root')
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()
cfg.use_gpu = torch.cuda.is_available()
if args.config_file:
cfg.merge_from_file(args.config_file)
reset_config(cfg, args)
cfg.merge_from_list(args.opts)
set_random_seed(cfg.train.seed)
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))
print('Collecting env info ...')
print('** System info **\n{}\n'.format(collect_env_info()))
if cfg.use_gpu:
torch.backends.cudnn.benchmark = True
datamanager = build_datamanager(cfg)
print('Building model: {}'.format(cfg.model.name))
model = torchreid.models.build_model(**model_kwargs(cfg, datamanager.num_train_pids))
num_params, flops = compute_model_complexity(model, (1, 3, cfg.data.height, cfg.data.width))
print('Model complexity: params={:,} flops={:,}'.format(num_params, flops))
if cfg.model.load_weights and check_isfile(cfg.model.load_weights):
if cfg.model.pretrained and not cfg.test.evaluate:
state_dict = torch.load(cfg.model.load_weights)
model.load_pretrained_weights(state_dict)
else:
load_pretrained_weights(model, cfg.model.load_weights)
if cfg.use_gpu:
model = nn.DataParallel(model).cuda()
optimizer = torchreid.optim.build_optimizer(model, **optimizer_kwargs(cfg))
scheduler = torchreid.optim.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, scheduler=scheduler
)
print('Building {}-engine for {}-reid'.format(cfg.loss.name, cfg.data.type))
engine = build_engine(cfg, datamanager, model, optimizer, scheduler)
engine.run(**engine_run_kwargs(cfg))
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() and not args.use_cpu)
log_name = 'test.log' if args.evaluate else 'train.log'
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))
torch.backends.cudnn.benchmark = True
else:
warnings.warn(
'Currently using CPU, however, GPU is highly recommended')
datamanager = build_datamanager(args)
print('Building model: {}'.format(args.arch))
model = torchreid.models.build_model(
name=args.arch,
num_classes=datamanager.num_train_pids,
loss=args.loss.lower(),
pretrained=(not args.no_pretrained),
use_gpu=use_gpu)
num_params, flops = compute_model_complexity(
model, (1, 3, args.height, args.width))
print('Model complexity: params={:,} flops={:,}'.format(num_params, flops))
if args.load_weights and check_isfile(args.load_weights):
load_pretrained_weights(model, args.load_weights)
if use_gpu:
model = nn.DataParallel(model).cuda()
optimizer = torchreid.optim.build_optimizer(model,
**optimizer_kwargs(args))
scheduler = torchreid.optim.build_lr_scheduler(optimizer,
**lr_scheduler_kwargs(args))
if args.resume and check_isfile(args.resume):
args.start_epoch = resume_from_checkpoint(args.resume,
model,
optimizer=optimizer)
print('Building {}-engine for {}-reid'.format(args.loss, args.app))
engine = build_engine(args, datamanager, model, optimizer, scheduler)
engine.run(**engine_run_kwargs(args))
def __init__(
self,
model_name='',
model_path='',
image_size=(256, 128),
pixel_mean=[0.485, 0.456, 0.406],
pixel_std=[0.229, 0.224, 0.225],
pixel_norm=True,
device='cuda',
verbose=True
):
# Build model
model = build_model(
model_name,
num_classes=1,
pretrained=True,
use_gpu=device.startswith('cuda')
)
model.eval()
num_params, flops = compute_model_complexity(
model, (1, 3, image_size[0], image_size[1])
)
if verbose:
print('Model: {}'.format(model_name))
print('- params: {:,}'.format(num_params))
print('- flops: {:,}'.format(flops))
if model_path and check_isfile(model_path):
load_pretrained_weights(model, model_path)
# Build transform functions
transforms = []
transforms += [T.Resize(image_size)]
transforms += [T.ToTensor()]
if pixel_norm:
transforms += [T.Normalize(mean=pixel_mean, std=pixel_std)]
preprocess = T.Compose(transforms)
to_pil = T.ToPILImage()
device = torch.device(device)
model.to(device)
# Class attributes
self.model = model
self.preprocess = preprocess
self.to_pil = to_pil
self.device = device
def main():
parser = argparse.ArgumentParser(
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('--model-name',
type=str,
default='',
help='Model name')
parser.add_argument('--weights', type=str, default='', help='Weights path')
parser.add_argument('--output',
type=str,
default='output',
help='Output path')
parser.add_argument('--resolution',
type=str,
default='128x256',
help='Resolution (WxH)')
args = parser.parse_args()
width, height = [int(i) for i in args.resolution.split('x')]
print('Collecting env info ...')
print('** System info **\n{}\n'.format(collect_env_info()))
imagedata_kwargs = {
'root': 'reid-data',
'sources': ['market1501'],
'targets': ['market1501'],
'height': 256,
'width': 128,
'transforms': ['random_flip', 'color_jitter'],
'norm_mean': [0.485, 0.456, 0.406],
'norm_std': [0.229, 0.224, 0.225],
'use_gpu': False,
'split_id': 0,
'combineall': False,
'load_train_targets': False,
'batch_size_train': 64,
'batch_size_test': 300,
'workers': 4,
'num_instances': 4,
'train_sampler': 'RandomSampler',
'cuhk03_labeled': False,
'cuhk03_classic_split': False,
'market1501_500k': False
}
datamanager = torchreid.data.ImageDataManager(**imagedata_kwargs)
print('Building model: {}'.format(args.model_name))
model = torchreid.models.build_model(
name=args.model_name,
num_classes=datamanager.num_train_pids,
loss='softmax',
pretrained=True,
use_gpu=False)
num_params, flops = compute_model_complexity(model, (1, 3, height, width))
print('Model complexity: params={:,} flops={:,}'.format(num_params, flops))
if args.weights and check_isfile(args.weights):
load_pretrained_weights(model, args.weights)
_input = torch.Tensor(1, 3, height, width)
inputs = (_input, )
print('Converting PyTorch model to ONNX...')
tmp = tempfile.mktemp(suffix='.onnx')
torch.onnx._export(model, inputs, tmp, export_params=True)
onnx_model = onnx.load(tmp)
export_path = args.output
onnx.checker.check_model(onnx_model)
print('Prepare TF model...')
tf_rep = prepare(onnx_model, strict=False)
if path.exists(export_path):
shutil.rmtree(export_path)
with tf.Session() as persisted_sess:
print("load graph")
persisted_sess.graph.as_default()
tf.import_graph_def(tf_rep.graph.as_graph_def(), name='')
i_tensors = []
o_tensors = []
inputs = {}
outputs = {}
for i in tf_rep.inputs:
t = persisted_sess.graph.get_tensor_by_name(
tf_rep.tensor_dict[i].name)
i_tensors.append(t)
tensor_info = tf.saved_model.utils.build_tensor_info(t)
inputs[t.name.split(':')[0].lower()] = tensor_info
print('input tensor [name=%s, type=%s, shape=%s]' %
(t.name, t.dtype.name, t.shape.as_list()))
print('')
for i in tf_rep.outputs:
t = persisted_sess.graph.get_tensor_by_name(
tf_rep.tensor_dict[i].name)
o_tensors.append(t)
tensor_info = tf.saved_model.utils.build_tensor_info(t)
outputs[t.name.split(':')[0]] = tensor_info
print('output tensor [name=%s, type=%s, shape=%s]' %
(t.name, t.dtype.name, t.shape.as_list()))
feed_dict = {}
for i in i_tensors:
feed_dict[i] = np.random.rand(*i.shape.as_list()).astype(
i.dtype.name)
print('test run:')
res = persisted_sess.run(o_tensors, feed_dict=feed_dict)
print(res)
# print('INPUTS')
# print(inputs)
# print('OUTPUTS')
# print(outputs)
prediction_signature = (
tf.saved_model.signature_def_utils.build_signature_def(
inputs=inputs,
outputs=outputs,
method_name=tf.saved_model.signature_constants.
PREDICT_METHOD_NAME))
builder = tf.saved_model.builder.SavedModelBuilder(export_path)
builder.add_meta_graph_and_variables(
persisted_sess, [tf.saved_model.tag_constants.SERVING],
signature_def_map={
tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
prediction_signature
})
builder.save()
print('Model saved to %s' % export_path)
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(
'-s',
'--sources',
type=str,
nargs='+',
help='source datasets (delimited by space)'
)
parser.add_argument(
'-t',
'--targets',
type=str,
nargs='+',
help='target datasets (delimited by space)'
)
parser.add_argument(
'--transforms', type=str, nargs='+', help='data augmentation'
)
parser.add_argument(
'--root', type=str, default='', help='path to data root'
)
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()
cfg.use_gpu = torch.cuda.is_available()
if args.config_file:
cfg.merge_from_file(args.config_file)
reset_config(cfg, args)
cfg.merge_from_list(args.opts)
set_random_seed(cfg.train.seed)
if cfg.use_gpu and args.gpu_devices:
# if gpu_devices is not specified, all available gpus will be used
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))
print('Collecting env info ...')
print('** System info **\n{}\n'.format(collect_env_info()))
if cfg.use_gpu:
torch.backends.cudnn.benchmark = True
datamanager = torchreid.data.ImageDataManager(**imagedata_kwargs(cfg))
print('Building model: {}'.format(cfg.model.name))
model = osnet_models.build_model(
cfg.model.name, num_classes=datamanager.num_train_pids
)
num_params, flops = compute_model_complexity(
model, (1, 3, cfg.data.height, cfg.data.width)
)
print('Model complexity: params={:,} flops={:,}'.format(num_params, flops))
if cfg.use_gpu:
model = nn.DataParallel(model).cuda()
optimizer = torchreid.optim.build_optimizer(model, **optimizer_kwargs(cfg))
scheduler = torchreid.optim.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
)
print('Building NAS engine')
engine = ImageSoftmaxNASEngine(
datamanager,
model,
optimizer,
scheduler=scheduler,
use_gpu=cfg.use_gpu,
label_smooth=cfg.loss.softmax.label_smooth,
mc_iter=cfg.nas.mc_iter,
init_lmda=cfg.nas.init_lmda,
min_lmda=cfg.nas.min_lmda,
lmda_decay_step=cfg.nas.lmda_decay_step,
lmda_decay_rate=cfg.nas.lmda_decay_rate,
fixed_lmda=cfg.nas.fixed_lmda
)
engine.run(**engine_run_kwargs(cfg))
print('*** Display the found architecture ***')
if cfg.use_gpu:
model.module.build_child_graph()
else:
model.build_child_graph()
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('-s',
'--sources',
type=str,
nargs='+',
help='source datasets (delimited by space)')
parser.add_argument('-t',
'--targets',
type=str,
nargs='+',
help='target datasets (delimited by space)')
parser.add_argument('--transforms',
type=str,
nargs='+',
help='data augmentation')
parser.add_argument('--root',
type=str,
default='',
help='path to data root')
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()
cfg.use_gpu = torch.cuda.is_available()
if args.config_file:
cfg.merge_from_file(args.config_file)
reset_config(cfg, args)
cfg.merge_from_list(args.opts)
set_random_seed(cfg.train.seed)
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))
print('Collecting env info ...')
print('** System info **\n{}\n'.format(collect_env_info()))
if cfg.use_gpu:
torch.backends.cudnn.benchmark = True
datamanager = torchreid.data.ImageDataManager(**imagedata_kwargs(cfg))
print('Building model-1: {}'.format(cfg.model.name))
model1 = torchreid.models.build_model(
name=cfg.model.name,
num_classes=datamanager.num_train_pids,
loss=cfg.loss.name,
pretrained=cfg.model.pretrained,
use_gpu=cfg.use_gpu)
num_params, flops = compute_model_complexity(
model1, (1, 3, cfg.data.height, cfg.data.width))
print('Model complexity: params={:,} flops={:,}'.format(num_params, flops))
print('Copying model-1 to model-2')
model2 = copy.deepcopy(model1)
if cfg.model.load_weights1 and check_isfile(cfg.model.load_weights1):
load_pretrained_weights(model1, cfg.model.load_weights1)
if cfg.model.load_weights2 and check_isfile(cfg.model.load_weights2):
load_pretrained_weights(model2, cfg.model.load_weights2)
if cfg.use_gpu:
model1 = nn.DataParallel(model1).cuda()
model2 = nn.DataParallel(model2).cuda()
optimizer1 = torchreid.optim.build_optimizer(model1,
**optimizer_kwargs(cfg))
scheduler1 = torchreid.optim.build_lr_scheduler(optimizer1,
**lr_scheduler_kwargs(cfg))
optimizer2 = torchreid.optim.build_optimizer(model2,
**optimizer_kwargs(cfg))
scheduler2 = torchreid.optim.build_lr_scheduler(optimizer2,
**lr_scheduler_kwargs(cfg))
if cfg.model.resume1 and check_isfile(cfg.model.resume1):
cfg.train.start_epoch = resume_from_checkpoint(cfg.model.resume1,
model1,
optimizer=optimizer1,
scheduler=scheduler1)
if cfg.model.resume2 and check_isfile(cfg.model.resume2):
resume_from_checkpoint(cfg.model.resume2,
model2,
optimizer=optimizer2,
scheduler=scheduler2)
print('Building DML-engine for image-reid')
engine = ImageDMLEngine(datamanager,
model1,
optimizer1,
scheduler1,
model2,
optimizer2,
scheduler2,
margin=cfg.loss.triplet.margin,
weight_t=cfg.loss.triplet.weight_t,
weight_x=cfg.loss.triplet.weight_x,
weight_ml=cfg.loss.dml.weight_ml,
use_gpu=cfg.use_gpu,
label_smooth=cfg.loss.softmax.label_smooth,
deploy=cfg.model.deploy)
engine.run(**engine_run_kwargs(cfg))
def main():
global args
set_random_seed(args.seed)
use_gpu = torch.cuda.is_available() and not args.use_cpu
log_name = 'test.log' if args.evaluate else 'train.log'
sys.stdout = Logger(osp.join(args.save_dir, log_name))
print('** Arguments **')
arg_keys = list(args.__dict__.keys())
arg_keys.sort()
for key in arg_keys:
print('{}: {}'.format(key, args.__dict__[key]))
print('\n')
print('Collecting env info ...')
print('** System info **\n{}\n'.format(collect_env_info()))
if use_gpu:
torch.backends.cudnn.benchmark = True
else:
warnings.warn(
'Currently using CPU, however, GPU is highly recommended')
dataset_vars = init_dataset(use_gpu)
trainloader, valloader, testloader, num_attrs, attr_dict = dataset_vars
if args.weighted_bce:
print('Use weighted binary cross entropy')
print('Computing the weights ...')
bce_weights = torch.zeros(num_attrs, dtype=torch.float)
for _, attrs, _ in trainloader:
bce_weights += attrs.sum(0) # sum along the batch dim
bce_weights /= len(trainloader) * args.batch_size
print('Sample ratio for each attribute: {}'.format(bce_weights))
bce_weights = torch.exp(-1 * bce_weights)
print('BCE weights: {}'.format(bce_weights))
bce_weights = bce_weights.expand(args.batch_size, num_attrs)
criterion = nn.BCEWithLogitsLoss(weight=bce_weights)
else:
print('Use plain binary cross entropy')
criterion = nn.BCEWithLogitsLoss()
print('Building model: {}'.format(args.arch))
model = models.build_model(args.arch,
num_attrs,
pretrained=not args.no_pretrained,
use_gpu=use_gpu)
num_params, flops = compute_model_complexity(
model, (1, 3, args.height, args.width))
print('Model complexity: params={:,} flops={:,}'.format(num_params, flops))
if args.load_weights and check_isfile(args.load_weights):
load_pretrained_weights(model, args.load_weights)
if use_gpu:
model = nn.DataParallel(model).cuda()
criterion = criterion.cuda()
if args.evaluate:
test(model, testloader, attr_dict, use_gpu)
return
optimizer = torchreid.optim.build_optimizer(model,
**optimizer_kwargs(args))
scheduler = torchreid.optim.build_lr_scheduler(optimizer,
**lr_scheduler_kwargs(args))
start_epoch = args.start_epoch
best_result = -np.inf
if args.resume and check_isfile(args.resume):
checkpoint = torch.load(args.resume)
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
start_epoch = checkpoint['epoch']
best_result = checkpoint['label_mA']
print('Loaded checkpoint from "{}"'.format(args.resume))
print('- start epoch: {}'.format(start_epoch))
print('- label_mA: {}'.format(best_result))
time_start = time.time()
for epoch in range(start_epoch, args.max_epoch):
train(epoch, model, criterion, optimizer, scheduler, trainloader,
use_gpu)
test_outputs = test(model, testloader, attr_dict, use_gpu)
label_mA = test_outputs[0]
is_best = label_mA > best_result
if is_best:
best_result = label_mA
save_checkpoint(
{
'state_dict': model.state_dict(),
'epoch': epoch + 1,
'label_mA': label_mA,
'optimizer': optimizer.state_dict(),
},
args.save_dir,
is_best=is_best)
elapsed = round(time.time() - time_start)
elapsed = str(datetime.timedelta(seconds=elapsed))
print('Elapsed {}'.format(elapsed))
def main():
# Load model configuration
parser = argparse.ArgumentParser()
parser.add_argument('-c',
'--config',
required=True,
help='path to configuration file')
args = parser.parse_args()
with open(args.config, "r") as ymlfile:
config = yaml.load(ymlfile, Loader=yaml.FullLoader)
# Automatically add sub-folder name to config["save_dir"], with the same name
# as the config file. For example, config["save_dir"] is typically "logs",
# so this would change config["save_dir"] to "logs/exp01", for example, so that
# we don't need to change the save_dir in every single config file (it instead
# automatically generates it from the name of the config file).
experiment_number = pathlib.Path(args.config).stem
config["save_dir"] = os.path.join(config["save_dir"], experiment_number)
# Set random seeds
set_random_seed(config["seed"])
# Set up GPU
if not config["use_avai_gpus"]:
os.environ['CUDA_VISIBLE_DEVICES'] = config["gpu_devices"]
use_gpu = torch.cuda.is_available() and not config["use_cpu"]
# Set up log files
log_name = 'test.log' if config["evaluate"] else 'train.log'
log_name += time.strftime('-%Y-%m-%d-%H-%M-%S')
sys.stdout = Logger(osp.join(config["save_dir"], log_name))
# Prepare for training
print('==========\nArgs:{}\n=========='.format(config))
print('Collecting env info ...')
print('** System info **\n{}\n'.format(collect_env_info()))
if use_gpu:
torch.backends.cudnn.benchmark = True
else:
warnings.warn(
'Currently using CPU, however, GPU is highly recommended')
# Build datamanager and model
datamanager = build_datamanager(config)
print('Building model: {}'.format(config["arch"]))
model = torchreid.models.build_model(
name=config["arch"],
num_classes=datamanager.num_train_pids,
loss=config["loss"].lower(),
pretrained=(not config["no_pretrained"]),
use_gpu=use_gpu)
# Compute model complexity
num_params, flops = compute_model_complexity(
model, (1, 3, config["height"], config["width"]))
print('Model complexity: params={:,} flops={:,}'.format(num_params, flops))
# Load pretrained weights if necessary
if config["load_weights"] and check_isfile(config["load_weights"]):
load_pretrained_weights(model, config["load_weights"])
# Set up multi-gpu
if use_gpu:
model = nn.DataParallel(model).cuda()
# Model settings
optimizer = torchreid.optim.build_optimizer(model,
**optimizer_kwargs(config))
scheduler = torchreid.optim.build_lr_scheduler(
optimizer, **lr_scheduler_kwargs(config))
if config["resume"] and check_isfile(config["resume"]):
config["start_epoch"] = resume_from_checkpoint(config["resume"],
model,
optimizer=optimizer)
print('Building {}-engine for {}-reid'.format(config["loss"],
config["app"]))
engine = build_engine(config, datamanager, model, optimizer, scheduler)
engine.run(**engine_run_kwargs(config))
