def __init__(self,
weights_path='./models/reid300.pt',
config_path='./models/reid300.yaml'):
if not os.path.exists(config_path):
raise ValueError("Invalid config path `" +
os.path.abspath(config_path) + "`")
if not os.path.exists(weights_path):
raise ValueError("Invalid weight path `" +
os.path.abspath(weights_path) + "`")
cfg = get_default_config()
cfg.merge_from_file(config_path)
self.model = build_model(
name=cfg.model.name,
num_classes=1000,
loss='am_softmax',
feature_dim=cfg.model.feature_dim,
fpn_cfg=cfg.model.fpn,
pooling_type=cfg.model.pooling_type,
input_size=(256, 128),
IN_first=cfg.model.IN_first,
)
load_pretrained_weights(self.model, weights_path)
self.model = self.model.cuda()
self.model.eval()
self.preprocess_mean = (torch.tensor([0.485, 0.456, 0.406]) *
255.0).cuda()
self.preprocess_std = (torch.tensor([0.229, 0.224, 0.225]) *
255.0).cuda()
def make_nncf_changes_in_training(model, cfg, classes, command_line_cfg_opts):
if cfg.train.ema.enable:
raise RuntimeError('EMA model could not be used together with NNCF compression')
if cfg.lr_finder.enable:
raise RuntimeError('LR finder could not be used together with NNCF compression')
if cfg.model.resume:
raise NotImplementedError('Resuming NNCF training is not implemented yet')
if not cfg.model.load_weights:
raise RuntimeError('NNCF training should be started from a pre-trained model')
checkpoint_path = cfg.model.load_weights
checkpoint_dict = load_checkpoint(checkpoint_path, map_location='cpu')
if is_nncf_state(checkpoint_dict):
raise RuntimeError(f'The checkpoint is NNCF checkpoint at {checkpoint_path}')
logger.info(f'Loading weights from {checkpoint_path}')
load_pretrained_weights(model, pretrained_dict=checkpoint_dict)
datamanager_for_init = build_datamanager(cfg, classes)
compression_ctrl, model, nncf_metainfo = \
wrap_nncf_model(model, cfg, datamanager_for_init=datamanager_for_init)
logger.info(f'Received from wrapping nncf_metainfo =\n{pformat(nncf_metainfo)}')
# calculating initial LR for NNCF training
lr = None
initial_lr_from_checkpoint = checkpoint_dict.get('initial_lr')
is_initial_lr_set_from_opts = is_config_parameter_set_from_command_line(command_line_cfg_opts, 'train.lr')
lr = calculate_lr_for_nncf_training(cfg, initial_lr_from_checkpoint,
is_initial_lr_set_from_opts)
assert lr is not None
cfg.train.lr = lr
return compression_ctrl, model, cfg, lr, nncf_metainfo
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 init_pretrained_weights(model, key='mobilenetv3', **kwargs):
"""Initializes model with pretrained weights.
"""
from torchreid.utils import load_pretrained_weights
link_to_weights = pretrained_urls[key]
load_pretrained_weights(model, link_to_weights, key, **kwargs)
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 _load_model(self,
model: ModelEntity,
device: torch.device,
pretrained_dict: Optional[Dict] = None):
if model is not None:
# If a model has been trained and saved for the task already, create empty model and load weights here
if pretrained_dict is None:
buffer = io.BytesIO(model.get_data("weights.pth"))
model_data = torch.load(buffer,
map_location=torch.device('cpu'))
else:
model_data = pretrained_dict
model = self._create_model(self._cfg, from_scratch=True)
try:
load_pretrained_weights(model, pretrained_dict=model_data)
logger.info("Loaded model weights from Task Environment")
except BaseException as ex:
raise ValueError("Could not load the saved model. The model file structure is invalid.") \
from ex
else:
# If there is no trained model yet, create model with pretrained weights as defined in the model config
# file.
model = self._create_model(self._cfg, from_scratch=False)
logger.info(
"No trained model in project yet. Created new model with general-purpose pretrained weights."
)
return model.to(device)
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():
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('--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)
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 _create_model(self, config, from_scratch: bool = False):
"""
Creates a model, based on the configuration in config
:param config: deep-object-reid configuration from which the model has to be built
:param from_scratch: bool, if True does not load any weights
:return model: Model in training mode
"""
num_train_classes = len(self._labels)
model = torchreid.models.build_model(
**model_kwargs(config, num_train_classes))
if self._cfg.model.load_weights and not from_scratch:
load_pretrained_weights(model, self._cfg.model.load_weights)
return model
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 __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 = 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('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)
cfg.merge_from_list(args.opts)
model = torchreid.models.build_model(name=cfg.model.name,
num_classes=2,
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,
activation=cfg.model.activation,
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)
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=True,
export_params=True)
def main():
parser = argparse.ArgumentParser()
parser.add_argument(
'--root',
type=str,
default=
'/media/ddj2/ce611f70-968b-4316-9547-9bc9cf931d32/V20200108/zhejiang_train'
)
parser.add_argument('-d', '--dataset', type=str, default='rock_dataset')
parser.add_argument('-m', '--model', type=str, default='abd_resnet')
parser.add_argument('--weights', type=str)
parser.add_argument('--save-dir', type=str, default='log/resnet50_cam')
parser.add_argument('--height', type=int, default=672)
parser.add_argument('--width', type=int, default=672)
args = parser.parse_args()
os.environ["CUDA_VISIBLE_DEVICES"] = "3"
use_gpu = torch.cuda.is_available()
torchreid.data.register_image_dataset(
'rock_dataset', torchreid.data.datasets.image.rock_dataset.RockDataSet)
datamanager = torchreid.data.ImageDataManager(
root=args.root,
sources=args.dataset,
height=args.height,
width=args.width,
batch_size_train=4,
batch_size_test=4,
transforms=None,
train_sampler='SequentialSampler')
test_loader = datamanager.test_loader
model = torchreid.models.build_model(
name=args.model,
num_classes=datamanager.num_train_pids,
use_gpu=use_gpu)
if use_gpu:
model = model.cuda()
if args.weights and check_isfile(args.weights):
load_pretrained_weights(model, args.weights)
visactmap(model, test_loader, args.save_dir, args.width, args.height,
use_gpu)
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--root', type=str, default='')
parser.add_argument('-m', '--model', type=str, default='resnet50')
parser.add_argument(
'--weights',
type=str,
default=
'/media/ddj2/ce611f70-968b-4316-9547-9bc9cf931d32/remote_data/PycharmProjects/ABD-Net-master/model_best.pth.tar'
)
parser.add_argument('--save-dir', type=str, default='logs/resnet50')
parser.add_argument('--height', type=int, default=672)
parser.add_argument('--width', type=int, default=672)
args = parser.parse_args()
use_gpu = torch.cuda.is_available()
test_dir = '/media/ddj2/ce611f70-968b-4316-9547-9bc9cf931d32/测试集/ceshi/crop/浙江省温州苍南县西古庵早白垩世小平田组PM201(挑选3张泛化测试用）20200114'
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
test_dataset = datasets.ImageFolder(
test_dir, transforms.Compose([
transforms.ToTensor(),
normalize,
]))
test_loader = torch.utils.data.DataLoader(test_dataset,
batch_size=8,
shuffle=False,
num_workers=8,
pin_memory=True)
model = torchreid.models.build_model(name=args.model,
num_classes=70,
use_gpu=use_gpu)
if use_gpu:
model = model.cuda()
if args.weights and check_isfile(args.weights):
load_pretrained_weights(model, args.weights)
visactmap(model, test_loader, args.save_dir, args.width, args.height,
use_gpu)
def test_single(num_attrs=26):
global args, best_accu
args = parser.parse_args()
attr_num = attr_nums['pa100k']
# create model
model = models.build_model(args.arch,
num_attrs,
pretrained=not args.no_pretrained,
use_gpu=True)
load_pretrained_weights(model, args.load_weights)
model = model.cuda()
model.eval()
img_path = "/home/bavon/face_test/reid/c1.jpg"
output = get_attr(img_path, model)
print("output is:{}".format(output))
for it in range(attr_num):
if output[it] == 1:
print('{} '.format(description['pa100k'][it]))
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 main():
parser = argparse.ArgumentParser()
parser.add_argument('--root', type=str)
parser.add_argument('-d', '--dataset', type=str, default='market1501')
parser.add_argument('-m', '--model', type=str, default='osnet_x1_0')
parser.add_argument('--weights', type=str)
parser.add_argument('--save-dir', type=str, default='log')
parser.add_argument('--height', type=int, default=256)
parser.add_argument('--width', type=int, default=128)
args = parser.parse_args()
use_gpu = torch.cuda.is_available()
datamanager = torchreid.data.ImageDataManager(
root=args.root,
sources=args.dataset,
height=args.height,
width=args.width,
batch_size_train=100,
batch_size_test=100,
transforms=None,
train_sampler='SequentialSampler'
)
test_loader = datamanager.test_loader
model = torchreid.models.build_model(
name=args.model,
num_classes=datamanager.num_train_pids,
use_gpu=use_gpu
)
if use_gpu:
model = model.cuda()
if args.weights and check_isfile(args.weights):
load_pretrained_weights(model, args.weights)
visactmap(
model, test_loader, args.save_dir, args.width, args.height, use_gpu
)
def test_dis(num_attrs=26):
global args, best_accu
args = parser.parse_args()
attr_num = attr_nums['pa100k']
# create model
model = models.build_model(args.arch,
num_attrs,
pretrained=not args.no_pretrained,
use_gpu=True)
load_pretrained_weights(model, args.load_weights)
model = model.cuda()
model.eval()
img1_path = '/home/bavon/face_test/reid/yuexin.jpg'
img1_path = '/home/bavon/face_test/reid/d1.jpg'
img2_path = '/home/bavon/model/datasets/rap/RAP_dataset/CAM17_2014-02-20_20140220175154-20140220175854_tarid124_frame2893_line1.png'
img2_path = '/home/bavon/face_test/reid/yangyi.jpg'
img2_path = '/home/bavon/face_test/reid/yuexin3.jpg'
img2_path = '/home/bavon/face_test/reid/d2.jpg'
output1 = get_attr(img1_path, model)
output2 = get_attr(img2_path, model)
dist = np.linalg.norm(output1 - output2)
print("distance is:{}".format(dist))
def download_model(net,
model_name,
local_model_store_dir_path=os.path.join(
"~", ".torch", "models")):
"""
Load model state dictionary from a file with downloading it if necessary.
Parameters
----------
net : Module
Network in which weights are loaded.
model_name : str
Name of the model.
local_model_store_dir_path : str, default $TORCH_HOME/models
Location for keeping the model parameters.
ignore_extra : bool, default True
Whether to silently ignore parameters from the file that are not present in this Module.
"""
net = load_pretrained_weights(
model=net,
file_path=get_model_file(
model_name=model_name,
local_model_store_dir_path=local_model_store_dir_path))
return net
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():
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')
d = CrowdHuman(
root='/mnt/lustre/share/fengweitao',
meta_file=
'/mnt/lustre/share/fengweitao/crowd_human/annotation_train.odgt')
dl = torch.utils.data.DataLoader(d, batch_size=16, num_workers=4)
print(dl)
config_file = 'configs/im_osnet_x1_0_softmax_256x128_amsgrad_cosine.yaml'
cfg = get_default_config()
cfg.use_gpu = torch.cuda.is_available()
cfg.merge_from_file(config_file)
model = torchreid.models.build_model(name=cfg.model.name,
num_classes=1024,
loss=cfg.loss.name,
pretrained=False,
use_gpu=cfg.use_gpu)
load_pretrained_weights(model, cfg.model.load_weights)
# m = None
model.eval()
model.cuda()
all_results = []
for i, data in enumerate(dl):
data = to_cuda(data, device='cuda')[0]
print(data.keys())
n, m, c, h, w = data['im'].shape
indata = data['im'].reshape(n * m, c, h, w)
with torch.no_grad():
o = model(indata)
o = o.reshape(n, m, -1)
dm = 0.
for j in range(o.size(1)):
for k in range(o.size(1)):
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('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():
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 train(self,
dataset: DatasetEntity,
output_model: ModelEntity,
train_parameters: Optional[TrainParameters] = None):
""" Trains a model on a dataset """
train_model = deepcopy(self._model)
if train_parameters is not None:
update_progress_callback = train_parameters.update_progress
else:
update_progress_callback = default_progress_callback
time_monitor = TrainingProgressCallback(
update_progress_callback,
num_epoch=self._cfg.train.max_epoch,
num_train_steps=math.ceil(
len(dataset.get_subset(Subset.TRAINING)) /
self._cfg.train.batch_size),
num_val_steps=0,
num_test_steps=0)
self.metrics_monitor = DefaultMetricsMonitor()
self.stop_callback.reset()
set_random_seed(self._cfg.train.seed)
train_subset = dataset.get_subset(Subset.TRAINING)
val_subset = dataset.get_subset(Subset.VALIDATION)
self._cfg.custom_datasets.roots = [
OTEClassificationDataset(train_subset,
self._labels,
self._multilabel,
keep_empty_label=self._empty_label
in self._labels),
OTEClassificationDataset(val_subset,
self._labels,
self._multilabel,
keep_empty_label=self._empty_label
in self._labels)
]
datamanager = torchreid.data.ImageDataManager(
**imagedata_kwargs(self._cfg))
num_aux_models = len(self._cfg.mutual_learning.aux_configs)
if self._cfg.use_gpu:
main_device_ids = list(range(self.num_devices))
extra_device_ids = [main_device_ids for _ in range(num_aux_models)]
train_model = DataParallel(train_model,
device_ids=main_device_ids,
output_device=0).cuda(
main_device_ids[0])
else:
extra_device_ids = [None for _ in range(num_aux_models)]
optimizer = torchreid.optim.build_optimizer(
train_model, **optimizer_kwargs(self._cfg))
if self._cfg.lr_finder.enable:
scheduler = None
else:
scheduler = torchreid.optim.build_lr_scheduler(
optimizer,
num_iter=datamanager.num_iter,
**lr_scheduler_kwargs(self._cfg))
if self._cfg.lr_finder.enable:
_, train_model, optimizer, scheduler = \
run_lr_finder(self._cfg, datamanager, train_model, optimizer, scheduler, None,
rebuild_model=False, gpu_num=self.num_devices, split_models=False)
_, final_acc = run_training(self._cfg,
datamanager,
train_model,
optimizer,
scheduler,
extra_device_ids,
self._cfg.train.lr,
tb_writer=self.metrics_monitor,
perf_monitor=time_monitor,
stop_callback=self.stop_callback)
training_metrics = self._generate_training_metrics_group()
self.metrics_monitor.close()
if self.stop_callback.check_stop():
logger.info('Training cancelled.')
return
logger.info("Training finished.")
best_snap_path = os.path.join(self._scratch_space, 'best.pth')
if os.path.isfile(best_snap_path):
load_pretrained_weights(self._model, best_snap_path)
for filename in os.listdir(self._scratch_space):
match = re.match(r'best_(aux_model_[0-9]+\.pth)', filename)
if match:
aux_model_name = match.group(1)
best_aux_snap_path = os.path.join(self._scratch_space,
filename)
self._aux_model_snap_paths[aux_model_name] = best_aux_snap_path
self.save_model(output_model)
performance = Performance(score=ScoreMetric(value=final_acc,
name="accuracy"),
dashboard_metrics=training_metrics)
logger.info(f'FINAL MODEL PERFORMANCE {performance}')
output_model.performance = performance
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)
