def main():
parser = ArgumentParser(formatter_class=ArgumentDefaultsHelpFormatter)
parser.add_argument('--config-file', '-c', type=str, required=True)
parser.add_argument('--root', '-r', type=str, required=True)
parser.add_argument('--save-dir', type=str, default='log')
parser.add_argument('opts', default=None, nargs=REMAINDER)
args = parser.parse_args()
assert osp.exists(args.config_file)
assert osp.exists(args.root)
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)
if cfg.use_gpu:
torch.backends.cudnn.benchmark = True
data_loader, num_pids = prepare_data(cfg, mode='gallery')
print('Building model: {}'.format(cfg.model.name))
model = torchreid.models.build_model(**model_kwargs(cfg, num_pids))
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 = model.cuda()
visualize_activation_map(model, data_loader, args.save_dir, cfg.data.width,
cfg.data.height, cfg.use_gpu)
def build_auxiliary_model(config_file,
num_classes,
use_gpu,
device_ids,
num_iter,
lr=None,
nncf_aux_config_changes=None,
aux_config_opts=None,
aux_pretrained_dict=None):
aux_cfg = get_default_config()
aux_cfg.use_gpu = use_gpu
merge_from_files_with_base(aux_cfg, config_file)
if nncf_aux_config_changes:
print(
f'applying to aux config changes from NNCF aux config {nncf_aux_config_changes}'
)
if not isinstance(nncf_aux_config_changes, CfgNode):
nncf_aux_config_changes = CfgNode(nncf_aux_config_changes)
aux_cfg.merge_from_other_cfg(nncf_aux_config_changes)
if aux_config_opts:
print(f'applying to aux config changes from command line arguments, '
f'the changes are:\n{pformat(aux_config_opts)}')
aux_cfg.merge_from_list(aux_config_opts)
print(f'\nShow auxiliary configuration\n{aux_cfg}\n')
if lr is not None:
aux_cfg.train.lr = lr
print(f"setting learning rate from main model: {lr}")
model = torchreid.models.build_model(**model_kwargs(aux_cfg, num_classes))
optimizer = torchreid.optim.build_optimizer(model,
**optimizer_kwargs(aux_cfg))
scheduler = torchreid.optim.build_lr_scheduler(
optimizer=optimizer, num_iter=num_iter, **lr_scheduler_kwargs(aux_cfg))
if aux_cfg.model.resume and check_isfile(aux_cfg.model.resume):
aux_cfg.train.start_epoch = resume_from_checkpoint(
aux_cfg.model.resume,
model,
optimizer=optimizer,
scheduler=scheduler)
elif aux_pretrained_dict is not None:
load_pretrained_weights(model, pretrained_dict=aux_pretrained_dict)
elif aux_cfg.model.load_weights and check_isfile(
aux_cfg.model.load_weights):
load_pretrained_weights(model, aux_cfg.model.load_weights)
if aux_cfg.use_gpu:
assert device_ids is not None
if len(device_ids) > 1:
model = DataParallel(model, device_ids=device_ids,
output_device=0).cuda(device_ids[0])
else:
model = model.cuda(device_ids[0])
return model, optimizer, scheduler
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():
parser = argparse.ArgumentParser(
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
description=
'The script adds the default int8 quantization NNCF metainfo '
'to NNCF deep-object-reid checkpoints '
'that were trained when NNCF metainfo was not '
'stored in NNCF checkpoints')
parser.add_argument('--config-file',
type=str,
required=True,
help='path to config file')
parser.add_argument('--checkpoint',
type=str,
required=True,
help='path to the src checkpoint file')
parser.add_argument(
'--dst-folder',
type=str,
required=True,
help='path to the dst folder to store dst checkpoint file')
args = parser.parse_args()
cfg = get_default_config()
merge_from_files_with_base(cfg, args.config_file)
checkpoint = torch.load(args.checkpoint, map_location='cpu')
if not isinstance(checkpoint, dict):
raise RuntimeError(
'Wrong format of checkpoint -- it is not the result of deep-object-reid training'
)
if checkpoint.get('nncf_metainfo'):
raise RuntimeError(
f'Checkpoint {args.checkpoint} already has nncf_metainfo')
if not os.path.isdir(args.dst_folder):
raise RuntimeError(f'The dst folder {args.dst_folder} is NOT present')
# default nncf config
h, w = cfg.data.height, cfg.data.width
nncf_config_data = get_default_nncf_compression_config(h, w)
nncf_metainfo = {
'nncf_compression_enabled': True,
'nncf_config': nncf_config_data
}
checkpoint['nncf_metainfo'] = nncf_metainfo
res_path = save_checkpoint(checkpoint, args.dst_folder)
def __init__(self, task_environment: TaskEnvironment):
logger.info("Loading OTEClassificationTask.")
self._scratch_space = tempfile.mkdtemp(prefix="ote-cls-scratch-")
logger.info(f"Scratch space created at {self._scratch_space}")
self._task_environment = task_environment
if len(task_environment.get_labels(False)) == 1:
self._labels = task_environment.get_labels(include_empty=True)
else:
self._labels = task_environment.get_labels(include_empty=False)
self._empty_label = get_empty_label(task_environment.label_schema)
self._multilabel = len(task_environment.label_schema.get_groups(False)) > 1 and \
len(task_environment.label_schema.get_groups(False)) == \
len(task_environment.get_labels(include_empty=False))
self._hierarchical = False
if not self._multilabel and len(
task_environment.label_schema.get_groups(False)) > 1:
self._labels = get_leaf_labels(task_environment.label_schema)
self._hierarchical = True
template_file_path = task_environment.model_template.model_template_path
self._base_dir = os.path.abspath(os.path.dirname(template_file_path))
self._cfg = get_default_config()
self._patch_config(self._base_dir)
if self._multilabel:
assert self._cfg.model.type == 'multilabel', task_environment.model_template.model_template_path + \
' model template does not support multiclass classification'
else:
assert self._cfg.model.type == 'classification', task_environment.model_template.model_template_path + \
' model template does not support multilabel classification'
self.device = torch.device(
"cuda:0") if torch.cuda.device_count() else torch.device("cpu")
self._model = self._load_model(task_environment.model,
device=self.device)
self.stop_callback = StopCallback()
self.metrics_monitor = DefaultMetricsMonitor()
# Set default model attributes.
self._optimization_methods = []
self._precision = [ModelPrecision.FP32]
self._optimization_type = ModelOptimizationType.MO
def __init__(self,
config_path='',
model_path='',
device='cuda',
verbose=True):
# Build model
cfg = get_default_config()
merge_from_files_with_base(cfg, config_path)
cfg.use_gpu = device.startswith('cuda')
model = build_model(**model_kwargs(cfg, 1))
model.eval()
image_size = (cfg.data.height, cfg.data.width)
flops, num_params = get_model_complexity_info(
model, (3, image_size[0], image_size[1]),
as_strings=False,
verbose=False,
print_per_layer_stat=False)
if verbose:
print('Model: {}'.format(cfg.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()]
print(cfg.data.norm_mean, cfg.data.norm_std)
transforms += [
T.Normalize(mean=cfg.data.norm_mean, std=cfg.data.norm_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 = ArgumentParser(formatter_class=ArgumentDefaultsHelpFormatter)
parser.add_argument('--config', '-c', type=str, required=True)
parser.add_argument('opts', default=None, nargs=REMAINDER)
args = parser.parse_args()
assert osp.exists(args.config)
cfg = get_default_config()
cfg.use_gpu = torch.cuda.is_available()
merge_from_files_with_base(cfg, args.config)
cfg.merge_from_list(args.opts)
cfg.freeze()
model = torchreid.models.build_model(**model_kwargs(cfg, [0, 0]))
load_pretrained_weights(model, cfg.model.load_weights)
conv_layers = collect_conv_layers(model)
show_stat(conv_layers)
def _patch_config(self, base_dir: str):
self._cfg = get_default_config()
if self._multilabel:
config_file_path = os.path.join(base_dir,
'main_model_multilabel.yaml')
else:
config_file_path = os.path.join(base_dir, 'main_model.yaml')
merge_from_files_with_base(self._cfg, config_file_path)
self._cfg.use_gpu = torch.cuda.device_count() > 0
self.num_devices = 1 if self._cfg.use_gpu else 0
self._cfg.custom_datasets.types = [
'external_classification_wrapper',
'external_classification_wrapper'
]
self._cfg.custom_datasets.names = ['train', 'val']
self._cfg.custom_datasets.roots = [''] * 2
self._cfg.data.sources = ['train']
self._cfg.data.targets = ['val']
self._cfg.data.save_dir = self._scratch_space
self._cfg.test.test_before_train = False
self.num_classes = len(self._labels)
for i, conf in enumerate(self._cfg.mutual_learning.aux_configs):
if str(base_dir) not in conf:
self._cfg.mutual_learning.aux_configs[i] = os.path.join(
base_dir, conf)
self._cfg.train.lr = self._hyperparams.learning_parameters.learning_rate
self._cfg.train.batch_size = self._hyperparams.learning_parameters.batch_size
self._cfg.test.batch_size = max(
1, self._hyperparams.learning_parameters.batch_size // 2)
self._cfg.train.max_epoch = self._hyperparams.learning_parameters.max_num_epochs
self._cfg.lr_finder.enable = self._hyperparams.learning_parameters.enable_lr_finder
self._cfg.train.early_stopping = self._hyperparams.learning_parameters.enable_early_stopping
def main():
parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('--config-file', type=str, default='', required=True,
help='Path to config file')
parser.add_argument('--output-name', type=str, default='model',
help='Path to save ONNX model')
parser.add_argument('--num-classes', type=int, nargs='+', default=None)
parser.add_argument('--opset', type=int, default=11)
parser.add_argument('--verbose', action='store_true',
help='Verbose mode for onnx.export')
parser.add_argument('--disable-dyn-axes', default=False, action='store_true')
parser.add_argument('--export_ir', action='store_true')
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:
merge_from_files_with_base(cfg, args.config_file)
reset_config(cfg)
cfg.merge_from_list(args.opts)
compression_hyperparams = get_compression_hyperparams(cfg.model.load_weights)
is_nncf_used = compression_hyperparams['enable_quantization'] or compression_hyperparams['enable_pruning']
if is_nncf_used:
print(f'Using NNCF -- making NNCF changes in config')
cfg = make_nncf_changes_in_config(cfg,
compression_hyperparams['enable_quantization'],
compression_hyperparams['enable_pruning'],
args.opts)
cfg.train.mix_precision = False
cfg.freeze()
num_classes = parse_num_classes(source_datasets=cfg.data.sources,
classification=cfg.model.type == 'classification' or cfg.model.type == 'multilabel',
num_classes=args.num_classes,
snap_path=cfg.model.load_weights)
model = build_model(**model_kwargs(cfg, num_classes))
if cfg.model.load_weights:
load_pretrained_weights(model, cfg.model.load_weights)
else:
warnings.warn("No weights are passed through 'load_weights' parameter! "
"The model will be converted with random or pretrained weights", category=RuntimeWarning)
if 'tresnet' in cfg.model.name:
patch_InplaceAbn_forward()
if is_nncf_used:
print('Begin making NNCF changes in model')
model = make_nncf_changes_in_eval(model, cfg)
print('End making NNCF changes in model')
onnx_file_path = export_onnx(model=model.eval(),
cfg=cfg,
output_file_path=args.output_name,
disable_dyn_axes=args.disable_dyn_axes,
verbose=args.verbose,
opset=args.opset,
extra_check=True)
if args.export_ir:
input_shape = [1, 3, cfg.data.height, cfg.data.width]
export_ir(onnx_model_path=onnx_file_path,
norm_mean=cfg.data.norm_mean,
norm_std=cfg.data.norm_std,
input_shape=input_shape,
optimized_model_dir=os.path.dirname(os.path.abspath(onnx_file_path)),
data_type='FP32')
def main():
parser = build_base_argparser()
parser.add_argument('-e',
'--auxiliary-models-cfg',
type=str,
nargs='*',
default='',
help='path to extra config files')
parser.add_argument('--split-models',
action='store_true',
help='whether to split models on own gpu')
parser.add_argument('--enable_quantization',
action='store_true',
help='Enable NNCF quantization algorithm')
parser.add_argument('--enable_pruning',
action='store_true',
help='Enable NNCF pruning algorithm')
parser.add_argument(
'--aux-config-opts',
nargs='+',
default=None,
help='Modify aux config options using the command-line')
args = parser.parse_args()
cfg = get_default_config()
cfg.use_gpu = torch.cuda.is_available() and args.gpu_num > 0
if args.config_file:
merge_from_files_with_base(cfg, args.config_file)
reset_config(cfg, args)
cfg.merge_from_list(args.opts)
is_nncf_used = args.enable_quantization or args.enable_pruning
if is_nncf_used:
print(f'Using NNCF -- making NNCF changes in config')
cfg = make_nncf_changes_in_config(cfg, args.enable_quantization,
args.enable_pruning, args.opts)
set_random_seed(cfg.train.seed, cfg.train.deterministic)
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))
if cfg.use_gpu:
torch.backends.cudnn.benchmark = True
num_aux_models = len(cfg.mutual_learning.aux_configs)
datamanager = build_datamanager(cfg, args.classes)
num_train_classes = datamanager.num_train_pids
print('Building main model: {}'.format(cfg.model.name))
model = torchreid.models.build_model(
**model_kwargs(cfg, num_train_classes))
macs, num_params = get_model_complexity_info(
model, (3, cfg.data.height, cfg.data.width),
as_strings=False,
verbose=False,
print_per_layer_stat=False)
print('Main model complexity: params={:,} flops={:,}'.format(
num_params, macs * 2))
aux_lr = cfg.train.lr # placeholder, needed for aux models, may be filled by nncf part below
if is_nncf_used:
print('Begin making NNCF changes in model')
if cfg.use_gpu:
model.cuda()
compression_ctrl, model, cfg, aux_lr, nncf_metainfo = \
make_nncf_changes_in_training(model, cfg,
args.classes,
args.opts)
should_freeze_aux_models = True
print(f'should_freeze_aux_models = {should_freeze_aux_models}')
print('End making NNCF changes in model')
else:
compression_ctrl = None
should_freeze_aux_models = False
nncf_metainfo = None
# creating optimizer and scheduler -- it should be done after NNCF part, since
# NNCF could change some parameters
optimizer = torchreid.optim.build_optimizer(model, **optimizer_kwargs(cfg))
if cfg.lr_finder.enable and not cfg.model.resume:
scheduler = None
else:
scheduler = torchreid.optim.build_lr_scheduler(
optimizer=optimizer,
num_iter=datamanager.num_iter,
**lr_scheduler_kwargs(cfg))
# Loading model (and optimizer and scheduler in case of resuming training).
# Note that if NNCF is used, loading is done inside NNCF part, so loading here is not required.
if cfg.model.resume and check_isfile(
cfg.model.resume) and not is_nncf_used:
device_ = 'cuda' if cfg.use_gpu else 'cpu'
cfg.train.start_epoch = resume_from_checkpoint(cfg.model.resume,
model,
optimizer=optimizer,
scheduler=scheduler,
device=device_)
elif cfg.model.load_weights and not is_nncf_used:
load_pretrained_weights(model, cfg.model.load_weights)
if cfg.model.type == 'classification':
check_classification_classes(model,
datamanager,
args.classes,
test_only=cfg.test.evaluate)
model, extra_device_ids = put_main_model_on_the_device(
model, cfg.use_gpu, args.gpu_num, num_aux_models, args.split_models)
if cfg.lr_finder.enable and not cfg.test.evaluate and not cfg.model.resume:
aux_lr, model, optimizer, scheduler = run_lr_finder(
cfg,
datamanager,
model,
optimizer,
scheduler,
args.classes,
rebuild_model=True,
gpu_num=args.gpu_num,
split_models=args.split_models)
log_dir = cfg.data.tb_log_dir if cfg.data.tb_log_dir else cfg.data.save_dir
run_training(cfg,
datamanager,
model,
optimizer,
scheduler,
extra_device_ids,
aux_lr,
tb_writer=SummaryWriter(log_dir=log_dir),
should_freeze_aux_models=should_freeze_aux_models,
nncf_metainfo=nncf_metainfo,
compression_ctrl=compression_ctrl)
def main():
parser = build_base_argparser()
args = parser.parse_args()
cfg = get_default_config()
cfg.use_gpu = torch.cuda.is_available() and args.gpu_num > 0
if args.config_file:
merge_from_files_with_base(cfg, args.config_file)
reset_config(cfg, args)
cfg.merge_from_list(args.opts)
is_ie_model = cfg.model.load_weights.endswith('.xml')
if not is_ie_model:
compression_hyperparams = get_compression_hyperparams(
cfg.model.load_weights)
is_nncf_used = compression_hyperparams[
'enable_quantization'] or compression_hyperparams['enable_pruning']
if is_nncf_used:
print(f'Using NNCF -- making NNCF changes in config')
cfg = make_nncf_changes_in_config(
cfg, compression_hyperparams['enable_quantization'],
compression_hyperparams['enable_pruning'], args.opts)
else:
is_nncf_used = False
set_random_seed(cfg.train.seed)
log_name = 'test.log' + time.strftime('-%Y-%m-%d-%H-%M-%S')
sys.stdout = Logger(osp.join(cfg.data.save_dir, log_name))
datamanager = torchreid.data.ImageDataManager(filter_classes=args.classes,
**imagedata_kwargs(cfg))
num_classes = len(
datamanager.test_loader[cfg.data.targets[0]]['query'].dataset.classes)
cfg.train.ema.enable = False
if not is_ie_model:
model = torchreid.models.build_model(**model_kwargs(cfg, num_classes))
load_pretrained_weights(model, cfg.model.load_weights)
if is_nncf_used:
print('Begin making NNCF changes in model')
model = make_nncf_changes_in_eval(model, cfg)
print('End making NNCF changes in model')
if cfg.use_gpu:
num_devices = min(torch.cuda.device_count(), args.gpu_num)
main_device_ids = list(range(num_devices))
model = DataParallel(model,
device_ids=main_device_ids,
output_device=0).cuda(main_device_ids[0])
else:
from torchreid.utils.ie_tools import VectorCNN
from openvino.inference_engine import IECore
cfg.test.batch_size = 1
model = VectorCNN(IECore(),
cfg.model.load_weights,
'CPU',
switch_rb=True,
**model_kwargs(cfg, num_classes))
for _, dataloader in datamanager.test_loader.items():
dataloader['query'].dataset.transform.transforms = \
dataloader['query'].dataset.transform.transforms[:-2]
if cfg.model.type == 'classification':
check_classification_classes(model,
datamanager,
args.classes,
test_only=True)
engine = build_engine(cfg=cfg,
datamanager=datamanager,
model=model,
optimizer=None,
scheduler=None)
engine.test(0,
dist_metric=cfg.test.dist_metric,
normalize_feature=cfg.test.normalize_feature,
visrank=cfg.test.visrank,
visrank_topk=cfg.test.visrank_topk,
save_dir=cfg.data.save_dir,
use_metric_cuhk03=cfg.cuhk03.use_metric_cuhk03,
ranks=(1, 5, 10, 20),
rerank=cfg.test.rerank)
def main():
# parse arguments
parser = build_base_argparser()
parser.add_argument('-e',
'--auxiliary-models-cfg',
type=str,
nargs='*',
default='',
help='path to extra config files')
parser.add_argument('--split-models',
action='store_true',
help='whether to split models on own gpu')
parser.add_argument(
'--aux-config-opts',
nargs='+',
default=None,
help='Modify aux config options using the command-line')
parser.add_argument('--epochs',
default=10,
type=int,
help='amount of the epochs')
args = parser.parse_args()
cfg = get_default_config()
cfg.use_gpu = torch.cuda.is_available() and args.gpu_num > 0
if args.config_file:
merge_from_files_with_base(cfg, args.config_file)
reset_config(cfg, args)
cfg.merge_from_list(args.opts)
set_random_seed(cfg.train.seed, cfg.train.deterministic)
log_name = 'optuna.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))
if cfg.use_gpu:
torch.backends.cudnn.benchmark = True
sampler = TPESampler(n_startup_trials=5, seed=True)
study = optuna.create_study(study_name='classification task',
direction="maximize",
sampler=sampler)
objective_partial = partial(objective, cfg, args)
try:
start_time = time.time()
study.optimize(objective_partial,
n_trials=cfg.lr_finder.n_trials,
timeout=None)
elapsed = round(time.time() - start_time)
print(
f"--- optimization is finished: {datetime.timedelta(seconds=elapsed)} ---"
)
except KeyboardInterrupt:
finish_process(study)
else:
finish_process(study)
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',
help='Path to save ONNX model')
parser.add_argument('--opset', type=int, default=9)
parser.add_argument('--verbose',
default=False,
action='store_true',
help='Verbose mode for onnx.export')
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)
cfg.merge_from_list(args.opts)
cfg.freeze()
num_classes = parse_num_classes(cfg.data.sources)
model = build_model(**model_kwargs(cfg, num_classes))
load_pretrained_weights(model, cfg.model.load_weights)
model.eval()
transform = build_inference_transform(
cfg.data.height,
cfg.data.width,
norm_mean=cfg.data.norm_mean,
norm_std=cfg.data.norm_std,
)
input_img = random_image(cfg.data.height, cfg.data.width)
input_blob = transform(input_img).unsqueeze(0)
input_names = ['data']
output_names = ['reid_embedding']
dynamic_axes = {
'data': {
0: 'batch_size',
1: 'channels',
2: 'height',
3: 'width'
},
'reid_embedding': {
0: 'batch_size',
1: 'dim'
}
}
output_file_path = args.output_name
if not args.output_name.endswith('.onnx'):
output_file_path += '.onnx'
register_op("group_norm", group_norm_symbolic, "", args.opset)
with torch.no_grad():
torch.onnx.export(
model,
input_blob,
output_file_path,
verbose=args.verbose,
export_params=True,
input_names=input_names,
output_names=output_names,
dynamic_axes=dynamic_axes,
opset_version=args.opset,
operator_export_type=torch.onnx.OperatorExportTypes.ONNX)
net_from_onnx = onnx.load(output_file_path)
try:
onnx.checker.check_model(net_from_onnx)
print('ONNX check passed.')
except onnx.onnx_cpp2py_export.checker.ValidationError as ex:
print('ONNX check failed: {}.'.format(ex))
def main():
parser = ArgumentParser(formatter_class=ArgumentDefaultsHelpFormatter)
parser.add_argument('--config-file', '-c', type=str, required=True)
parser.add_argument('--weights', '-w', type=str, required=True)
parser.add_argument('--root', '-r', type=str, required=True)
parser.add_argument('--out-dir', '-o', type=str, required=True)
parser.add_argument('--matrix-size',
'-ms',
type=int,
required=False,
default=8)
parser.add_argument('opts', default=None, nargs=REMAINDER)
args = parser.parse_args()
assert exists(args.config_file)
assert exists(args.weights)
assert exists(args.root)
create_dirs(args.out_dir)
cfg = get_default_config()
cfg.use_gpu = torch.cuda.is_available()
if args.config_file:
merge_from_files_with_base(cfg, args.config_file)
reset_config(cfg, args)
cfg.merge_from_list(args.opts)
if cfg.use_gpu:
torch.backends.cudnn.benchmark = True
target_dataset = 'cityflow'
data_query = build_query(cfg, target_dataset)
data_gallery, gallery_size = build_gallery(cfg, target_dataset)
print('Building model: {}'.format(cfg.model.name))
model = torchreid.models.build_model(**model_kwargs(cfg, [0]))
print('Loading model: {}'.format(cfg.model.load_weights))
cfg.model.load_weights = args.weights
load_pretrained_weights(model, cfg.model.load_weights)
model = model.cuda() if cfg.use_gpu else model
print('Extracting query embeddings ...')
images_query, embeddings_query, ids_query = run_model(
model, data_query, cfg.use_gpu)
print('Extracting gallery embeddings ...')
images_gallery, embeddings_gallery, ids_gallery = run_model(
model, data_gallery, cfg.use_gpu)
print('Calculating distance matrices ...')
distance_matrix_qg = calculate_distances(embeddings_query,
embeddings_gallery)
print('Finding matches ...')
top_k = args.matrix_size**2 - 1
matches = find_matches(distance_matrix_qg, top_k=top_k)
print('Dumping visualizations ...')
visualize_matches(matches, images_query, images_gallery, ids_query,
ids_gallery, args.matrix_size, args.out_dir)
def main():
parser = argparse.ArgumentParser(
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('--config-file',
type=str,
default='',
required=True,
help='path to config file')
parser.add_argument(
'--custom-roots',
type=str,
nargs='+',
help=
'types or paths to annotation of custom datasets (delimited by space)')
parser.add_argument('--custom-types',
type=str,
nargs='+',
help='path of custom datasets (delimited by space)')
parser.add_argument('--custom-names',
type=str,
nargs='+',
help='names of custom datasets (delimited by space)')
parser.add_argument('--root',
type=str,
default='',
help='path to data root')
parser.add_argument('--classes',
type=str,
nargs='+',
help='name of classes in classification dataset')
parser.add_argument('--out')
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:
merge_from_files_with_base(cfg, args.config_file)
reset_config(cfg, args)
cfg.merge_from_list(args.opts)
set_random_seed(cfg.train.seed)
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, args.classes)
num_train_classes = datamanager.num_train_pids
print('Building main model: {}'.format(cfg.model.name))
model = torchreid.models.build_model(
**model_kwargs(cfg, num_train_classes))
macs, num_params = get_model_complexity_info(
model, (3, cfg.data.height, cfg.data.width),
as_strings=False,
verbose=False,
print_per_layer_stat=False)
print('Main model complexity: M params={:,} G flops={:,}'.format(
num_params / 10**6, macs * 2 / 10**9))
if args.out:
out = list()
out.append({
'key': 'size',
'display_name': 'Size',
'value': num_params / 10**6,
'unit': 'Mp'
})
out.append({
'key': 'complexity',
'display_name': 'Complexity',
'value': 2 * macs / 10**9,
'unit': 'GFLOPs'
})
print('dump to' + args.out)
with open(args.out, 'w') as write_file:
json.dump(out, write_file, indent=4)
