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('--output-name', type=str, default='model')
parser.add_argument('--verbose',
default=False,
action='store_true',
help='Verbose mode for onnx.export')
args = parser.parse_args()
cfg = get_default_config()
if args.config_file:
cfg.merge_from_file(args.config_file)
cfg.freeze()
model = build_model(
name=cfg.model.name,
num_classes=1041, # Does not matter in conversion
loss=cfg.loss.name,
pretrained=False,
use_gpu=True,
feature_dim=cfg.model.feature_dim,
fpn=cfg.model.fpn,
fpn_dim=cfg.model.fpn_dim,
gap_as_conv=cfg.model.gap_as_conv,
input_size=(cfg.data.height, cfg.data.width),
IN_first=cfg.model.IN_first)
load_pretrained_weights(model, cfg.model.load_weights)
model.eval()
_, transform = build_transforms(cfg.data.height,
cfg.data.width,
transforms=cfg.data.transforms,
norm_mean=cfg.data.norm_mean,
norm_std=cfg.data.norm_std,
apply_masks_to_test=False)
input_size = (cfg.data.height, cfg.data.width, 3)
img = np.random.rand(*input_size).astype(np.float32)
img = np.uint8(img * 255)
im = Image.fromarray(img)
blob = transform(im).unsqueeze(0)
torch.onnx.export(
model,
blob,
args.output_name + '.onnx',
verbose=False,
export_params=True,
input_names=['data'],
output_names=['reid_embedding'],
opset_version=9) # 9th version resolves nearest upsample issue
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)
cfg.freeze()
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 = 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,
dropout_prob=cfg.model.dropout_prob,
feature_dim=cfg.model.feature_dim,
fpn=cfg.model.fpn.enable,
fpn_dim=cfg.model.fpn.dim,
gap_as_conv=cfg.model.gap_as_conv,
input_size=(cfg.data.height, cfg.data.width),
IN_first=cfg.model.IN_first,
fpn_process=cfg.model.fpn.process,
)
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 = build_lr_scheduler(optimizer, **lr_scheduler_kwargs(cfg))
if cfg.model.resume and check_isfile(cfg.model.resume):
args.start_epoch = resume_from_checkpoint(cfg.model.resume,
model,
optimizer=optimizer)
if len(cfg.model.openvino.name):
openvino_model = OpenVINOModel(cfg.model.openvino.name,
cfg.model.openvino.cpu_extension)
else:
openvino_model = None
print('Building {}-engine for {}-reid'.format(cfg.loss.name,
cfg.data.type))
engine = build_engine(cfg,
datamanager,
model,
optimizer,
scheduler,
openvino_model=openvino_model)
engine.run(**engine_run_kwargs(cfg))
# data_manager
data_manager = build_data_manager(cfg)
model = osnet_x1_0(num_classes=data_manager.num_train_pids,
loss=cfg.loss.name,
pretrained=cfg.model.pretrained,
use_gpu=cfg.use_gpu)
if cfg.model.load_weights:
load_pretrained_weights(model, cfg.model.load_weights)
model = model.cuda()
optimizer = build_optimizer(model, **optimizer_kwargs(cfg))
scheduler = build_scheduler(optimizer, **lr_scheduler_kwargs(cfg))
# print(model)
engine = build_engine(cfg, data_manager, model, optimizer, scheduler)
engine.run(**engine_run_kwargs(cfg))
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument(
'--config-file',
type=str,
default='./config/im_osnet_x1_0_softmax_256x128_amsgrad_cosine.yaml')
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)
train(cfg)