def get_basic_sparsity_config(model_size=4,
input_sample_size=(1, 1, 4, 4),
sparsity_init=0.02,
sparsity_target=0.5,
sparsity_steps=2,
sparsity_training_steps=3):
config = Config()
config.update({
"model": "basic_sparse_conv",
"model_size": model_size,
"input_info": {
"sample_size": input_sample_size,
},
"compression": {
"algorithm": "rb_sparsity",
"params": {
"schedule": "polynomial",
"sparsity_init": sparsity_init,
"sparsity_target": sparsity_target,
"sparsity_steps": sparsity_steps,
"sparsity_training_steps": sparsity_training_steps
},
"layers": {
"conv": {
"sparsify": True
},
}
}
})
return config
def test_json_against_nncf_config_schema(config_test_struct):
config_path, should_pass = config_test_struct
if should_pass:
_ = Config.from_json(str(config_path))
else:
with pytest.raises(jsonschema.ValidationError):
_ = Config.from_json(str(config_path))
def get_empty_config(model_size=4, input_sample_size=(1, 1, 4, 4)):
config = Config()
config.update({
"model": "basic_sparse_conv",
"model_size": model_size,
"input_sample_size": input_sample_size,
})
return config
def __init__(self, algo: 'QuantizationController', config: Config,
default_activation_bitwidth: int, default_weight_bitwidth: int, criterion: _Loss,
data_loader: DataLoader, is_distributed: bool = False):
super().__init__(algo, config, default_activation_bitwidth, default_weight_bitwidth,
criterion, data_loader, is_distributed)
self._traces_per_layer_path = config.get('traces_per_layer_path', None)
self._num_data_points = config.get('num_data_points', 200)
self._iter_number = config.get('iter_number', 200)
self._tolerance = config.get('tolerance', 1e-5)
self._bits = config.get('bits', [4, 8])
def get_basic_magnitude_sparsity_config(input_sample_size=(1, 1, 4, 4)):
config = Config()
config.update({
"model": "basic_sparse_conv",
"input_sample_size": input_sample_size,
"compression":
{
"algorithm": "magnitude_sparsity",
"params": {}
}
})
return config
def get_basic_sparsity_plus_quantization_config(input_sample_size=(1, 1, 4,
4)):
config = Config()
config.update({
"input_sample_size":
input_sample_size,
"compression": [{
"algorithm": "rb_sparsity",
}, {
"algorithm": "quantization",
}]
})
return config
def get_squeezenet_quantization_config(model_size=32):
config = Config()
config.update({
"model": "squeezenet1_1_custom",
"model_size": model_size,
"input_sample_size": (3, 3, model_size, model_size),
"compression": {
"algorithm": "quantization",
"initializer": {
"num_init_steps": 0
}
}
})
return config
def get_basic_quantization_config(model_size=4):
config = Config()
config.update({
"model": "basic_quant_conv",
"model_size": model_size,
"input_sample_size": (1, 1, model_size, model_size),
"compression": {
"algorithm": "quantization",
"initializer": {
"num_init_steps": 0
},
"params": {}
}
})
return config
def get_basic_pruning_config(input_sample_size=(1, 1, 4, 4)):
config = Config()
config.update({
"model": "pruning_conv_model",
"input_info":
{
"sample_size": input_sample_size,
},
"compression":
{
"params": {
}
}
})
return config
def __init__(self, pruning_algo, params: Config = None):
super().__init__()
if params is None:
self._params = Config()
else:
self._params = params
self.algo = pruning_algo
# Number of initial steps of training before pruning
self.num_init_steps = self._params.get('num_init_steps', 0)
self.pruning_steps = self._params.get('pruning_steps', 100)
# Pruning rates
self.initial_pruning = self._params.get('pruning_init', 0)
self.pruning_target = self._params.get('pruning_target', 0.5)
def __init__(self, algo: 'QuantizationController', config: Config,
default_activation_bitwidth: int, default_weight_bitwidth: int,
criterion: _Loss, data_loader: DataLoader, is_distributed: bool = False):
self._algo = algo
self._model = self._algo._model # type: NNCFNetwork
self._bitwidth_per_scope = config.get('bitwidth_per_scope', {}) # type: List[List]
self._default_activation_bitwidth = default_activation_bitwidth
self._default_weight_bitwidth = default_weight_bitwidth
self._criterion = criterion
self._data_loader = data_loader
self._is_distributed = is_distributed
self._all_quantizations = {}
self._ordered_weight_quantizations = []
for class_type in QUANTIZATION_MODULES.registry_dict.values():
quantization_type = class_type.__name__
act_module_dict = self._model.get_compression_modules_by_type(
CompressionModuleType.ACTIVATION_QUANTIZER)
func_module_dict = self._model.get_compression_modules_by_type(CompressionModuleType.FUNCTION_QUANTIZER)
weight_module_dict = self._model.get_nncf_wrapped_model()
self._all_quantizations.update(get_all_modules_by_type(act_module_dict, quantization_type))
self._all_quantizations.update(get_all_modules_by_type(func_module_dict, quantization_type))
ops_quantizations = get_all_modules_by_type(weight_module_dict, quantization_type)
self._ordered_weight_quantizations.extend([q for q in ops_quantizations.values() if q.is_weights])
self._all_quantizations.update(ops_quantizations)
def main(argv):
parser = get_argument_parser()
args = parser.parse_args(args=argv)
config = Config.from_json(args.config)
config.update_from_args(args, parser)
if config.dist_url == "env://":
config.update_from_env()
configure_paths(config)
source_root = Path(__file__).absolute().parents[2] # nncf root
create_code_snapshot(source_root,
osp.join(config.log_dir, "snapshot.tar.gz"))
if config.seed is not None:
warnings.warn('You have chosen to seed training. '
'This will turn on the CUDNN deterministic setting, '
'which can slow down your training considerably! '
'You may see unexpected behavior when restarting '
'from checkpoints.')
config.execution_mode = get_execution_mode(config)
if not is_binarization(config):
start_worker(main_worker, config)
else:
from examples.classification.binarization_worker import main_worker_binarization
start_worker(main_worker_binarization, config)
def start_evaluation(args):
"""Launches the evaluation process"""
if args.dataset == 'vgg':
dataset = VGGFace2(args.val,
args.v_list,
args.v_land,
landmarks_training=True)
elif args.dataset == 'celeb':
dataset = CelebA(args.val, args.v_land, test=True)
else:
dataset = NDG(args.val, args.v_land)
if dataset.have_landmarks:
log.info('Use alignment for the train data')
dataset.transform = t.Compose(
[Rescale((48, 48)), ToTensor(switch_rb=True)])
else:
exit()
val_loader = DataLoader(dataset,
batch_size=args.val_batch_size,
num_workers=4,
shuffle=False,
pin_memory=True)
model = models_landmarks['landnet']()
assert args.snapshot is not None
if args.compr_config:
config = Config.from_json(args.compr_config)
compression_algo = create_compression_algorithm(model, config)
model = compression_algo.model
log.info('Testing snapshot ' + args.snapshot + ' ...')
model = load_model_state(model,
args.snapshot,
args.device,
eval_state=True)
model.eval()
cudnn.benchmark = True
model = torch.nn.DataParallel(
model,
device_ids=[args.device],
)
log.info('Face landmarks model:')
log.info(model)
avg_err, per_point_avg_err, failures_rate = evaluate(val_loader, model)
log.info('Avg RMSE error: {}'.format(avg_err))
log.info('Per landmark RMSE error: {}'.format(per_point_avg_err))
log.info('Failure rate: {}'.format(failures_rate))
if args.compr_config and "sparsity_level" in compression_algo.statistics():
log.info("Sparsity level: {0:.2f}".format(
compression_algo.statistics()
['sparsity_rate_for_sparsified_modules']))
def __init__(self, binarization_algo, config=None):
super().__init__()
if config is None:
config = Config()
c = config['params']
self.config = config
self.algo = binarization_algo
self.activations_bin_start_epoch = c.get('activations_bin_start_epoch',
1)
self.weights_bin_start_epoch = c.get('weights_bin_start_epoch', 1)
self._set_binarization_status()
def ssd_vgg300():
ssd_params = Config({
"clip": False,
"variance": [0.1, 0.1, 0.2, 0.2],
"max_sizes": [60, 111, 162, 213, 264, 315],
"min_sizes": [30, 60, 111, 162, 213, 264],
"steps": [8, 16, 32, 64, 100, 300],
"aspect_ratios": [[2], [2, 3], [2, 3], [2, 3], [2], [2]],
"flip": True
})
return SSD_VGG(ssd_params, 300, 21, True)
class PruningScheduler(CompressionScheduler):
def __init__(self, pruning_algo, params: Config = None):
super().__init__()
if params is None:
self._params = Config()
else:
self._params = params
self.algo = pruning_algo
# Number of initial steps of training before pruning
self.num_init_steps = self._params.get('num_init_steps', 0)
self.pruning_steps = self._params.get('pruning_steps', 100)
# Pruning rates
self.initial_pruning = self._params.get('pruning_init', 0)
self.pruning_target = self._params.get('pruning_target', 0.5)
def load_state_dict(self, state_dict):
super().load_state_dict(state_dict)
self._set_pruning_level()
def epoch_step(self, epoch=None):
super().epoch_step(epoch)
self._set_pruning_level()
def _set_pruning_level(self):
self.algo.set_pruning_rate(self.current_pruning_level)
if self.last_epoch >= (self.pruning_steps + self.num_init_steps):
self.algo.freeze()
def _calc_density_level(self):
raise NotImplementedError
@property
def current_pruning_level(self):
if self.last_epoch >= self.num_init_steps:
return 1 - self._calc_density_level()
return 0
def ssd_mobilenet():
ssd_params = Config({
"variance": [0.1, 0.1, 0.2, 0.2],
"max_sizes": [60, 111, 162, 213, 264, 315],
"min_sizes": [30, 60, 111, 162, 213, 264],
"steps": [16, 32, 64, 100, 150, 300],
"aspect_ratios": [[2], [2, 3], [2, 3], [2, 3], [2], [2]],
"clip": False,
"flip": True,
"top_k": 200
})
return MobileNetSSD(21, ssd_params)
def main(argv):
parser = get_argument_parser()
args = parser.parse_args(args=argv)
config = Config.from_json(args.config)
config.update_from_args(args, parser)
configure_paths(config)
source_root = Path(__file__).absolute().parents[2] # nncf root
create_code_snapshot(source_root, osp.join(config.log_dir, "snapshot.tar.gz"))
config.execution_mode = get_execution_mode(config)
if config.dataset_dir is not None:
config.train_imgs = config.train_anno = config.test_imgs = config.test_anno = config.dataset_dir
start_worker(main_worker, config)
def test_scheduler_can_do_epoch_step__with_rb_algo():
config = Config()
config['input_info'] = [{"sample_size": [1, 1, 32, 32]}]
config['compression']['algorithm'] = 'rb_sparsity'
config['compression']["params"] = {
'schedule': 'polynomial',
'power': 1,
'sparsity_steps': 2,
'sparsity_init': 0.2,
'sparsity_target': 0.6,
'sparsity_training_steps': 4
}
_, compression_ctrl = create_compressed_model_and_algo_for_test(
BasicConvTestModel(), config)
scheduler = compression_ctrl.scheduler
loss = compression_ctrl.loss
assert pytest.approx(loss.target_sparsity_rate) == 0.2
assert not loss.disabled
for module_info in compression_ctrl.sparsified_module_info:
assert module_info.operand.sparsify
scheduler.epoch_step()
assert pytest.approx(loss.target_sparsity_rate, abs=1e-3) == 0.4
assert pytest.approx(loss().item(), abs=1e-3) == 64
assert not loss.disabled
scheduler.epoch_step()
assert pytest.approx(loss.target_sparsity_rate, abs=1e-3) == 0.6
assert pytest.approx(loss().item(), abs=1e-3) == 144
assert not loss.disabled
scheduler.epoch_step()
assert not loss.disabled
assert loss.target_sparsity_rate == 0.6
assert loss().item() == 144
scheduler.epoch_step()
assert loss.disabled
assert loss.target_sparsity_rate == 0.6
assert loss() == 0
for module_info in compression_ctrl.sparsified_module_info:
assert not module_info.operand.sparsify
def main(argv):
parser = get_common_argument_parser()
arguments = parser.parse_args(args=argv)
config = Config.from_json(arguments.config)
config.update_from_args(arguments, parser)
if config.dist_url == "env://":
config.update_from_env()
if config.mode.lower() != 'test':
if not osp.exists(config.log_dir):
os.makedirs(config.log_dir)
config.log_dir = str(config.log_dir)
configure_paths(config)
print("Save directory:", config.log_dir)
else:
config.log_dir = "/tmp/"
config.execution_mode = get_execution_mode(config)
start_worker(main_worker, config)
def train(args):
"""Launches training of landmark regression model"""
input_size = models_landmarks['landnet']().get_input_res()
if args.dataset == 'vgg':
drops_schedule = [1, 6, 9, 13]
dataset = VGGFace2(args.train, args.t_list, args.t_land, landmarks_training=True)
elif args.dataset == 'celeba':
drops_schedule = [10, 20]
dataset = CelebA(args.train, args.t_land)
else:
drops_schedule = [90, 140, 200]
dataset = NDG(args.train, args.t_land)
if dataset.have_landmarks:
log.info('Use alignment for the train data')
dataset.transform = transforms.Compose([landmarks_augmentation.Rescale((56, 56)),
landmarks_augmentation.Blur(k=3, p=.2),
landmarks_augmentation.HorizontalFlip(p=.5),
landmarks_augmentation.RandomRotate(50),
landmarks_augmentation.RandomScale(.8, .9, p=.4),
landmarks_augmentation.RandomCrop(48),
landmarks_augmentation.ToTensor(switch_rb=True)])
else:
log.info('Error: training dataset has no landmarks data')
exit()
train_loader = DataLoader(dataset, batch_size=args.train_batch_size, num_workers=4, shuffle=True)
writer = SummaryWriter('./logs_landm/{:%Y_%m_%d_%H_%M}_'.format(datetime.datetime.now()) + args.snap_prefix)
model = models_landmarks['landnet']()
set_dropout_fn = model.set_dropout_ratio
compression_algo = None
if args.snap_to_resume is not None:
config = Config.from_json(args.compr_config)
compression_algo = create_compression_algorithm(model, config)
model = compression_algo.model
log.info('Resuming snapshot ' + args.snap_to_resume + ' ...')
model = load_model_state(model, args.snap_to_resume, args.device, eval_state=False)
model = torch.nn.DataParallel(model, device_ids=[args.device])
def main():
model_bin, model_xml = get_ir_paths(args.model, args.bin)
config = Config.from_json(args.config)
input_infos_list = create_input_infos(config)
image_size = input_infos_list[0].shape[-1]
size = int(image_size / 0.875)
print('IE version: {}'.format(get_version()))
# NOTE: importing torch after loading IE to plugin to avoid issue with built-in MKLDNN of PyTorch
plugin = IEPlugin(device='CPU',
plugin_dirs=args.cpu_plugin_dir)
plugin.add_cpu_extension(os.path.join(args.cpu_plugin_dir, "libcpu_extension.so"))
net = IENetwork(model=model_xml, weights=model_bin)
exec_net = getExecNet(plugin, net)
from torch.utils.data import DataLoader
import torchvision.datasets as datasets
import torchvision.transforms as transforms
val_loader = DataLoader(
datasets.ImageFolder(args.data, transforms.Compose([
transforms.Resize(size),
transforms.CenterCrop(image_size),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])),
batch_size=1, shuffle=False, num_workers=4, pin_memory=True)
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir)
config['log_dir'] = args.output_dir
infer_fn = partial(infer_ie_model, net=net)
validate_general(val_loader, exec_net, infer_fn)
validate_torch_model(os.path.join(args.output_dir, "PTH"), config=config, num_layers=args.num_layers,
dump=args.dump, val_loader=val_loader, cuda=args.cuda)
def test_model_can_be_loaded_with_resume(_params, tmp_path):
p = _params
config_path = p['nncf_config_path']
checkpoint_path = p['checkpoint_path']
config = Config.from_json(str(config_path))
config.execution_mode = p['execution_mode']
config.current_gpu = 0
config.log_dir = str(tmp_path)
config.device = get_device(config)
config.distributed = config.execution_mode in (ExecutionMode.DISTRIBUTED, ExecutionMode.MULTIPROCESSING_DISTRIBUTED)
if config.distributed:
config.dist_url = "tcp://127.0.0.1:9898"
config.dist_backend = "nccl"
config.rank = 0
config.world_size = 1
configure_distributed(config)
model_name = config['model']
model = load_model(model_name,
pretrained=False,
num_classes=config.get('num_classes', 1000),
model_params=config.get('model_params'))
patch_torch_operators()
compression_algo, model = create_compressed_model(model, config)
model, _ = prepare_model_for_execution(model, config)
if config.distributed:
compression_algo.distributed()
reset_context('orig')
reset_context('quantized_graphs')
checkpoint = torch.load(checkpoint_path, map_location='cpu')
load_state(model, checkpoint['state_dict'], is_resume=True)
def test_get_default_weight_decay(algo, ref_weight_decay):
config = Config()
config.update({"compression": {"algorithm": algo}})
assert ref_weight_decay == get_default_weight_decay(config)
def train(args):
"""Performs training of a face recognition network"""
input_size = models_backbones[args.model]().get_input_res()
if args.train_dataset == 'vgg':
assert args.t_list
dataset = VGGFace2(args.train, args.t_list, args.t_land)
elif args.train_dataset == 'imdbface':
dataset = IMDBFace(args.train, args.t_list)
elif args.train_dataset == 'trp':
dataset = TrillionPairs(args.train, args.t_list)
else:
dataset = MSCeleb1M(args.train, args.t_list)
if dataset.have_landmarks:
log.info('Use alignment for the train data')
dataset.transform = t.Compose([
augm.HorizontalFlipNumpy(p=.5),
augm.CutOutWithPrior(p=0.05, max_area=0.1),
augm.RandomRotationNumpy(10, p=.95),
augm.ResizeNumpy(input_size),
augm.BlurNumpy(k=5, p=.2),
augm.NumpyToTensor(switch_rb=True)
])
else:
dataset.transform = t.Compose([
augm.ResizeNumpy(input_size),
augm.HorizontalFlipNumpy(),
augm.RandomRotationNumpy(10),
augm.NumpyToTensor(switch_rb=True)
])
if args.weighted:
train_weights = dataset.get_weights()
train_weights = torch.DoubleTensor(train_weights)
sampler = torch.utils.data.sampler.WeightedRandomSampler(
train_weights, len(train_weights))
train_loader = torch.utils.data.DataLoader(
dataset,
batch_size=args.train_batch_size,
sampler=sampler,
num_workers=3,
pin_memory=False)
else:
train_loader = DataLoader(dataset,
batch_size=args.train_batch_size,
num_workers=4,
shuffle=True)
lfw = LFW(args.val, args.v_list, args.v_land)
if lfw.use_landmarks:
log.info('Use alignment for the test data')
lfw.transform = t.Compose(
[augm.ResizeNumpy(input_size),
augm.NumpyToTensor(switch_rb=True)])
else:
lfw.transform = t.Compose([
augm.ResizeNumpy((160, 160)),
augm.CenterCropNumpy(input_size),
augm.NumpyToTensor(switch_rb=True)
])
log_path = './logs/{:%Y_%m_%d_%H_%M}_{}'.format(datetime.datetime.now(),
args.snap_prefix)
writer = SummaryWriter(log_path)
if not osp.exists(args.snap_folder):
os.mkdir(args.snap_folder)
model = models_backbones[args.model](embedding_size=args.embed_size,
num_classes=dataset.get_num_classes(),
feature=False)
set_dropout_fn = model.set_dropout_ratio
compression_algo = None
if args.snap_to_resume is not None:
if args.compr_config:
config = Config.from_json(args.compr_config)
compression_algo = create_compression_algorithm(model, config)
model = compression_algo.model
log.info('Resuming snapshot ' + args.snap_to_resume + ' ...')
model = load_model_state(model,
args.snap_to_resume,
args.devices[0],
eval_state=False)
model = torch.nn.DataParallel(model, device_ids=args.devices)
else:
model = torch.nn.DataParallel(model,
device_ids=args.devices,
output_device=args.devices[0])
model.cuda()
model.train()
cudnn.benchmark = True
if args.to_onnx is not None:
if args.compr_config:
compression_algo.export_model(args.to_onnx)
else:
model = model.eval().cpu()
input_shape = tuple([1, 3] + list(input_size))
with torch.no_grad():
torch.onnx.export(model.module,
torch.randn(input_shape),
args.to_onnx,
verbose=True)
print("Saved to", args.to_onnx)
return
log.info('Face Recognition model:')
log.info(model)
if args.mining_type == 'focal':
softmax_criterion = AMSoftmaxLoss(gamma=args.gamma,
m=args.m,
margin_type=args.margin_type,
s=args.s)
else:
softmax_criterion = AMSoftmaxLoss(t=args.t,
m=0.35,
margin_type=args.margin_type,
s=args.s)
aux_losses = MetricLosses(dataset.get_num_classes(), args.embed_size,
writer)
optimizer = optim.SGD(model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
if args.compr_config:
scheduler = optim.lr_scheduler.MultiStepLR(optimizer, [0, 2, 4, 6, 8])
else:
scheduler = optim.lr_scheduler.MultiStepLR(optimizer, [3, 6, 9, 13])
log.info('Epoch length: %d' % len(train_loader))
for epoch_num in range(args.epoch_total_num):
log.info('Epoch: %d' % epoch_num)
scheduler.step()
if epoch_num > 6 or args.compr_config:
set_dropout_fn(0.)
classification_correct = 0
classification_total = 0
for i, data in enumerate(train_loader, 0):
iteration = epoch_num * len(train_loader) + i
if iteration % args.val_step == 0:
snapshot_name = osp.join(
args.snap_folder,
args.snap_prefix + '_{0}.pt'.format(iteration))
if iteration > 0:
log.info('Saving Snapshot: ' + snapshot_name)
save_model_cpu(model, optimizer, snapshot_name, epoch_num)
log.info('Evaluating Snapshot: ' + snapshot_name)
model.eval()
same_acc, diff_acc, all_acc, auc = evaluate(
args,
lfw,
model,
compute_embeddings_lfw,
args.val_batch_size,
verbose=False)
model.train()
log.info('Validation accuracy: {0:.4f}, {1:.4f}'.format(
same_acc, diff_acc))
log.info('Validation accuracy mean: {0:.4f}'.format(all_acc))
log.info('Validation AUC: {0:.4f}'.format(auc))
writer.add_scalar('Epoch', epoch_num, iteration)
writer.add_scalar('Accuracy/Val_same_accuracy', same_acc,
iteration)
writer.add_scalar('Accuracy/Val_diff_accuracy', diff_acc,
iteration)
writer.add_scalar('Accuracy/Val_accuracy', all_acc, iteration)
writer.add_scalar('Accuracy/AUC', auc, iteration)
data, label = data['img'], data['label'].cuda()
features, sm_outputs = model(data)
optimizer.zero_grad()
aux_losses.init_iteration()
aux_loss, aux_log = aux_losses(features, label, epoch_num,
iteration)
loss_sm = softmax_criterion(sm_outputs, label)
compr_loss = compression_algo.loss() if args.compr_config else 0
loss = loss_sm + aux_loss + compr_loss
loss.backward()
aux_losses.end_iteration()
optimizer.step()
_, predicted = torch.max(sm_outputs.data, 1)
classification_total += int(label.size(0))
classification_correct += int(torch.sum(predicted.eq(label)))
train_acc = float(classification_correct) / classification_total
if i % 10 == 0:
log.info('Iteration %d, Softmax loss: %.4f, Total loss: %.4f' %
(iteration, loss_sm, loss) + aux_log)
log.info('Learning rate: %f' % scheduler.get_lr()[0])
writer.add_scalar('Loss/train_loss', loss, iteration)
writer.add_scalar('Loss/softmax_loss', loss_sm, iteration)
writer.add_scalar('Learning_rate',
scheduler.get_lr()[0], iteration)
writer.add_scalar('Accuracy/classification', train_acc,
iteration)
if args.compr_config and "sparsity_level" in compression_algo.statistics(
):
log.info('Sparsity_level: %.4f' %
compression_algo.statistics()["sparsity_level"])
writer.add_scalar(
'Sparsity_level',
compression_algo.statistics()["sparsity_level"],
iteration)
if args.compr_config:
compression_algo.scheduler.step()
if args.compr_config:
compression_algo.scheduler.epoch_step()
def main():
parser = argparse.ArgumentParser(description='Evaluation script for Face Recognition in PyTorch')
parser.add_argument('--devices', type=int, nargs='+', default=[0], help='CUDA devices to use.')
parser.add_argument('--embed_size', type=int, default=128, help='Size of the face embedding.')
parser.add_argument('--val_data_root', dest='val', required=True, type=str, help='Path to validation data.')
parser.add_argument('--val_list', dest='v_list', required=True, type=str, help='Path to train data image list.')
parser.add_argument('--val_landmarks', dest='v_land', default='', required=False, type=str,
help='Path to landmarks for the test images.')
parser.add_argument('--val_batch_size', type=int, default=8, help='Validation batch size.')
parser.add_argument('--snap', type=str, required=False, help='Snapshot to evaluate.')
parser.add_argument('--roc_fname', type=str, default='', help='ROC file.')
parser.add_argument('--dump_embeddings', action='store_true', help='Dump embeddings to summary writer.')
parser.add_argument('--dist', choices=['l2', 'cos'], type=str, default='cos', help='Distance.')
parser.add_argument('--flipped_emb', action='store_true', help='Flipped embedding concatenation trick.')
parser.add_argument('--show_failed', action='store_true', help='Show misclassified pairs.')
parser.add_argument('--model', choices=models_backbones.keys(), type=str, default='rmnet', help='Model type.')
parser.add_argument('--engine', choices=['pt', 'ie'], type=str, default='pt', help='Framework to use for eval.')
# IE-related options
parser.add_argument('--fr_model', type=str, required=False)
parser.add_argument('--lm_model', type=str, required=False)
parser.add_argument('-pp', '--plugin_dir', type=str, default=None, help='Path to a plugin folder')
parser.add_argument('-c', '--compr_config', help='Path to a file with compression parameters', required=False)
args = parser.parse_args()
if args.engine == 'pt':
assert args.snap is not None, 'To evaluate PyTorch snapshot, please, specify --snap option.'
if args.compr_config:
patch_torch_operators()
with torch.cuda.device(args.devices[0]):
data, embeddings_fun = load_test_dataset(args)
model = models_backbones[args.model](embedding_size=args.embed_size, feature=True)
if args.compr_config:
config = Config.from_json(args.compr_config)
compression_algo = create_compression_algorithm(model, config)
model = compression_algo.model
model = load_model_state(model, args.snap, args.devices[0])
evaluate(args, data, model, embeddings_fun, args.val_batch_size, args.dump_embeddings,
args.roc_fname, args.snap, True, args.show_failed)
if args.compr_config and "sparsity_level" in compression_algo.statistics():
log.info("Sparsity level: {0:.2f}".format(
compression_algo.statistics()['sparsity_rate_for_sparsified_modules']))
else:
from utils.ie_tools import load_ie_model
assert args.fr_model is not None, 'To evaluate IE model, please, specify --fr_model option.'
fr_model = load_ie_model(args.fr_model, 'CPU', args.plugin_dir)
lm_model = None
if args.lm_model:
lm_model = load_ie_model(args.lm_model, 'CPU', args.plugin_dir)
input_size = tuple(fr_model.get_input_shape()[2:])
lfw = LFW(args.val, args.v_list, args.v_land)
if not lfw.use_landmarks or lm_model:
lfw.transform = t.Compose([ResizeNumpy(220), CenterCropNumpy(input_size)])
lfw.use_landmarks = False
else:
log.info('Using landmarks for the LFW images.')
lfw.transform = t.Compose([ResizeNumpy(input_size)])
evaluate(args, lfw, fr_model, partial(compute_embeddings_lfw_ie, lm_model=lm_model), val_batch_size=1,
dump_embeddings=False, roc_fname='', snap_name='', verbose=True, show_failed=False)
def create_compressed_model(
model: Module,
config: Config,
dummy_forward_fn: Callable[[Module], Any] = None,
dump_graphs=True
) -> Tuple[CompressionAlgorithmController, NNCFNetwork]:
"""
The main function used to produce a model ready for compression fine-tuning from an original PyTorch
model and a configuration object.
dummy_forward_fn
:param model: The original model. Should have its parameters already loaded from a checkpoint or another
source.
:param config: A configuration object used to determine the exact compression modifications to be applied
to the model
:param dummy_forward_fn: will be used instead of a *forward* function call to build
the internal graph representation via tracing. Specifying this is useful when the original training pipeline
has special formats of data loader output or has additional *forward* arguments other than input tensors.
Otherwise, the *forward* call of the model during graph tracing will be made with mock tensors according
to the shape specified in the config object.
:param dump_graphs: Whether or not should also dump the internal graph representation of the
original and compressed models in the .dot format into the log directory.
:return: A controller for the compression algorithm (or algorithms, in which case the controller
is an instance of CompositeCompressionController) and the model ready for compression wrapped
as an object of NNCFNetwork."""
if dump_graphs:
if dummy_forward_fn is None:
input_info_list = create_input_infos(config)
graph_builder = GraphBuilder(
custom_forward_fn=create_dummy_forward_fn(
input_info_list, with_input_tracing=True))
else:
graph_builder = GraphBuilder(custom_forward_fn=dummy_forward_fn)
if is_main_process():
graph = graph_builder.build_graph(model)
graph.dump_graph(osp.join(config.log_dir, "original_graph.dot"),
extended=True)
if is_debug():
set_debug_log_dir(config.log_dir)
input_info_list = create_input_infos(config)
scopes_without_shape_matching = config.get('scopes_without_shape_matching',
[])
ignored_scopes = config.get('ignored_scopes')
target_scopes = config.get('target_scopes')
compressed_model = NNCFNetwork(
model,
input_infos=input_info_list,
dummy_forward_fn=dummy_forward_fn,
ignored_scopes=ignored_scopes,
target_scopes=target_scopes,
scopes_without_shape_matching=scopes_without_shape_matching)
compression_algo_builder_list = create_compression_algorithm_builders(
config)
for builder in compression_algo_builder_list:
compressed_model = builder.apply_to(compressed_model)
compression_ctrl = compressed_model.commit_compression_changes()
if dump_graphs and is_main_process() and compression_algo_builder_list:
if dummy_forward_fn is None:
compressed_graph_builder = GraphBuilder(
custom_forward_fn=create_dummy_forward_fn(
input_info_list, with_input_tracing=False))
else:
compressed_graph_builder = GraphBuilder(
custom_forward_fn=dummy_forward_fn)
graph = compressed_graph_builder.build_graph(
compressed_model, compressed_model.get_tracing_context())
graph.dump_graph(osp.join(config.log_dir, "compressed_graph.dot"),
extended=True)
return compression_ctrl, compressed_model
