def test_can_restore_binary_mask_on_magnitude_algo_resume(): config = get_empty_config() config['compression'] = { "algorithm": "magnitude_sparsity", "weight_importance": "abs", "params": { "schedule": "multistep", "sparsity_levels": [0.3, 0.5] } } magnitude_algo = create_compression_algorithm(MagnitudeTestModel(), config) sparse_model = magnitude_algo.model with torch.no_grad(): sparse_model(torch.ones([1, 1, 10, 10])) config = get_empty_config() config["compression"] = {"algorithm": "const_sparsity"} const_algo = create_compression_algorithm(MagnitudeTestModel(), config) const_sparse_model = const_algo.model load_state(const_sparse_model, sparse_model.state_dict()) op = const_sparse_model.conv1.pre_ops['0'] check_equal(ref_mask_1, op.operand.binary_mask) op = const_sparse_model.conv2.pre_ops['0'] check_equal(ref_mask_2, op.operand.binary_mask)
def test_load_state_interoperability(_algos, _model_wrapper, is_resume): config_save = get_empty_config() config_save['compression'] = [{ 'algorithm': algo, 'params': {} } for algo in _algos['save_algos']] algo_save = create_test_compression_algo(config_save, BasicConvTestModel()) model_save = _model_wrapper['save_model'](algo_save.model) saved_model_state = model_save.state_dict() ref_num_loaded = len(saved_model_state) config_resume = get_empty_config() config_resume['compression'] = [{ 'algorithm': algo, 'params': {} } for algo in _algos['load_algos']] algo_resume = create_test_compression_algo(config_resume, BasicConvTestModel()) model_resume = _model_wrapper['resume_model'](algo_resume.model) if not is_resume or (is_resume and _algos['is_resume_ok']): act_num_loaded = load_state(model_resume, saved_model_state, is_resume) if ('magnitude_sparsity' in _algos['load_algos'] or 'const_sparsity' in _algos['load_algos']) \ and 'rb_sparsity' in _algos['save_algos']: # no need to load _mask and _uniform ref_num_loaded -= 2 assert act_num_loaded == ref_num_loaded else: with pytest.raises(RuntimeError): load_state(model_resume, saved_model_state, is_resume)
def create_model(config): input_info_list = create_input_infos(config) image_size = input_info_list[0].shape[-1] ssd_net = build_ssd(config.model, config.ssd_params, image_size, config.num_classes, config) compression_algo, ssd_net = create_compressed_model(ssd_net, config) ssd_net = compression_algo.model weights = config.get('weights') if weights: sd = torch.load(weights, map_location='cpu') load_state(ssd_net, sd) ssd_net.train() model, _ = prepare_model_for_execution(ssd_net, config) return compression_algo, model
def build_ssd_vgg(cfg, size, num_classes, config): ssd_vgg = SSD_VGG(cfg, size, num_classes, batch_norm=config.get('batchnorm', False)) if config.basenet and (config.resuming_checkpoint is None) and (config.weights is None): print('Loading base network...') basenet_weights = torch.load(config.basenet) new_weights = {} for wn, wv in basenet_weights.items(): wn = wn.replace('features.', '') new_weights[wn] = wv load_state(ssd_vgg.basenet, new_weights, is_resume=False) return ssd_vgg
def build_ssd_mobilenet(cfg, size, num_classes, config): if size != 300: raise ValueError("Only Mobilenet-SSD with input size 300 is supported") mobilenet_ssd = MobileNetSSD(num_classes, cfg) if config.basenet and (config.resuming_checkpoint is None) and (config.weights is None): print('Loading base network...') basenet_weights = torch.load(config.basenet)['state_dict'] new_weights = {} for wn, wv in basenet_weights.items(): wn = wn.replace('model.', '') new_weights[wn] = wv load_state(mobilenet_ssd.basenet, new_weights, is_resume=False) return mobilenet_ssd
def build_ssd_vgg(cfg, size, num_classes, config): ssd_vgg = SSD_VGG(cfg, size, num_classes, batch_norm=config.get('batchnorm', False)) print('Initializing weights...') # ssd_vgg.apply(weights_init) if config.basenet: print('Loading base network...') basenet_weights = torch.load(config.basenet) new_weights = {} for wn, wv in basenet_weights.items(): wn = wn.replace('features.', '') new_weights[wn] = wv load_state(ssd_vgg.basenet, new_weights, is_resume=False) return ssd_vgg
def test_can_restore_binary_mask_on_magnitude_quant_algo_resume(tmp_path): patch_torch_operators() config = get_empty_config() config["compression"] = [{ "algorithm": "magnitude_sparsity", "weight_importance": "abs", "params": { "schedule": "multistep", "sparsity_levels": [0.3, 0.5] } }, { "algorithm": "quantization" }] config.log_dir = str(tmp_path) reset_context('orig') reset_context('quantized_graphs') _, model = create_compressed_model(MagnitudeTestModel(), config) # load_state doesn't support CPU + Quantization sparse_model = torch.nn.DataParallel(model) sparse_model.cuda() with torch.no_grad(): sparse_model(torch.ones([1, 1, 10, 10])) reset_context('orig') reset_context('quantized_graphs') config = get_empty_config() config.log_dir = str(tmp_path) config["compression"] = [{ "algorithm": "const_sparsity" }, { "algorithm": "quantization" }] _, const_sparse_model = create_compressed_model(MagnitudeTestModel(), config) load_state(const_sparse_model, sparse_model.state_dict()) op = const_sparse_model.get_nncf_wrapped_module().conv1.pre_ops['0'] check_equal(ref_mask_1, op.operand.binary_mask) op = const_sparse_model.get_nncf_wrapped_module().conv2.pre_ops['0'] check_equal(ref_mask_2, op.operand.binary_mask)
def resume_from_checkpoint(resuming_checkpoint, model, config, optimizer, compression_algo): best_acc1 = 0 if osp.isfile(resuming_checkpoint): print("=> loading checkpoint '{}'".format(resuming_checkpoint)) checkpoint = torch.load(resuming_checkpoint, map_location='cpu') load_state(model, checkpoint['state_dict'], is_resume=True) if config.mode.lower() == 'train' and config.to_onnx is None: config.start_epoch = checkpoint['epoch'] best_acc1 = checkpoint['best_acc1'] compression_algo.scheduler.load_state_dict(checkpoint['scheduler']) optimizer.load_state_dict(checkpoint['optimizer']) print("=> loaded checkpoint '{}' (epoch: {}, best_acc1: {:.3f})". format(resuming_checkpoint, checkpoint['epoch'], best_acc1)) else: print("=> loaded checkpoint '{}'".format(resuming_checkpoint)) else: raise FileNotFoundError( "no checkpoint found at '{}'".format(resuming_checkpoint)) return model, config, optimizer, compression_algo, best_acc1
def load_checkpoint(model, model_path, device_name, optimizer=None, compression_scheduler=None): """Loads the model from a specified directory with a specified name Keyword arguments: - model (``nn.Module``): The stored model state is copied to this model instance. - model_path: The model filename. - device_name: Device name for the model to be loaded into. - is_ddp: If true, model will be treated as a DistributedDataParallel instance and the actual model will be loaded into model.module - optimizer (``torch.optim``): The stored optimizer state is copied to this optimizer instance. - compression_algo: The compression scheduler for the saved state to be loaded into Returns: The ``model``, ``optimizer``, epoch, mean IoU and ``compression_scheduler``, loaded from the checkpoint. """ assert os.path.isfile( model_path), "The model file \"{0}\" doesn't exist.".format(model_path) # Load the stored model parameters to the model instance checkpoint = torch.load(model_path, map_location=device_name) load_state(model, checkpoint['state_dict'], is_resume=True) if optimizer is not None: optimizer.load_state_dict(checkpoint['optimizer']) epoch = checkpoint['epoch'] miou = checkpoint['miou'] if "scheduler" in checkpoint and compression_scheduler is not None: compression_scheduler.load_state_dict(checkpoint['scheduler']) return model, optimizer, epoch, miou, compression_scheduler
def test_load_state_skips_not_matched_params__from_smaller_to_larger(): ref_weights = torch.tensor([[[[3, 2], [2, 3]]]]) ref_bias = torch.tensor([2.]) model_save = BasicConvTestModel(out_channels=2) model_load = BasicConvTestModel(out_channels=1, weight_init=2, bias_init=2) num_loaded = load_state(model_load, model_save.state_dict()) assert num_loaded == 0 act_bias = model_load.conv.bias.data act_weights = model_load.conv.weight.data check_equal(act_bias, ref_bias) check_equal(act_weights, ref_weights)
def test_load_state_sets_initialized_flag(tmp_path): config = get_basic_quantization_config() config.log_dir = str(tmp_path) reset_context('orig') reset_context('quantized_graphs') _, model = create_compressed_model(TwoConvTestModel(), config) load_state( model, { 'module.features.0.0.pre_ops.0.op.signed_tensor': torch.tensor( [1.0]), # quantizer of 1st conv's weights 'module.features.1.0.pre_ops.0.op.scale': torch.tensor( [1.0]) # quantizer of 2nd conv's weights }) quantizers = get_all_modules_by_type(model, 'SymmetricQuantizer') for name, module in quantizers.items(): if 'activation_quantizers' in name or 'UpdateInputs' in name: assert not module.initialized else: assert module.initialized
def test_load_state_skips_not_matched_params__from_larger_to_smaller(): ref_weights = BasicConvTestModel.default_weight() ref_bias = BasicConvTestModel.default_bias() model_save = BasicConvTestModel(out_channels=1, weight_init=2, bias_init=2) model_load = BasicConvTestModel(out_channels=2) num_loaded = load_state(model_load, model_save.state_dict()) act_bias = model_load.conv.bias.data act_weights = model_load.conv.weight.data assert num_loaded == 0 check_equal(act_bias, ref_bias) check_equal(act_weights, ref_weights)
def test_scale_and_sign_init_for_quant_algo__after_load_state( self, wrap_dataloader): config = self.create_config() algo = self.create_algo(config) load_state( algo.model, { 'module.features.0.0.pre_ops.0.op.signed_tensor': torch.tensor([0.]), # quantizer of 1st conv's weights 'module.features.1.0.pre_ops.0.op.scale': torch.tensor([100]) # quantizer of 2nd conv's weights }) data_loader = self.create_dataloader(wrap_dataloader, config, algo) algo.initialize(data_loader) self.check_sign_and_scale( algo, { '.*Sequential\\[0\\].*UpdateWeight.*': (False, 1), '.*Sequential\\[1\\].*UpdateWeight. *': (False, 100), '.*activation_quantizers.*Sequential\\[0\\].*': (True, 4), '.*activation_quantizers.*Sequential\\[1\\].*': (True, 24) })
def test_ordinary_load(algo, _model_wrapper, is_resume): config = get_empty_config() if algo: config['compression'] = {'algorithm': algo, 'params': {}} algo_save = create_test_compression_algo(config, BasicConvTestModel()) model_save = _model_wrapper['save_model'](algo_save.model) algo_resume = create_test_compression_algo(config, BasicConvTestModel()) model_resume = _model_wrapper['resume_model'](algo_resume.model) num_loaded = load_state(model_resume, model_save.state_dict(), is_resume) assert num_loaded == len(model_save.state_dict())
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 main_worker(current_gpu, config): config.current_gpu = current_gpu config.distributed = config.execution_mode in ( ExecutionMode.DISTRIBUTED, ExecutionMode.MULTIPROCESSING_DISTRIBUTED) if config.distributed: configure_distributed(config) config.device = get_device(config) if is_main_process(): configure_logging(config) print_args(config) if config.seed is not None: manual_seed(config.seed) cudnn.deterministic = True cudnn.benchmark = False # create model model_name = config['model'] weights = config.get('weights') model = load_model(model_name, pretrained=config.get('pretrained', True) if weights is None else False, num_classes=config.get('num_classes', 1000), model_params=config.get('model_params')) compression_algo, model = create_compressed_model(model, config) if weights: load_state(model, torch.load(weights, map_location='cpu')) model, _ = prepare_model_for_execution(model, config) if config.distributed: compression_algo.distributed() is_inception = 'inception' in model_name # define loss function (criterion) and optimizer criterion = nn.CrossEntropyLoss() criterion = criterion.to(config.device) params_to_optimize = get_parameter_groups(model, config) optimizer, lr_scheduler = make_optimizer(params_to_optimize, config) resuming_checkpoint = config.resuming_checkpoint best_acc1 = 0 # optionally resume from a checkpoint if resuming_checkpoint is not None: model, config, optimizer, compression_algo, best_acc1 = \ resume_from_checkpoint(resuming_checkpoint, model, config, optimizer, compression_algo) if config.to_onnx is not None: compression_algo.export_model(config.to_onnx) print("Saved to", config.to_onnx) return if config.execution_mode != ExecutionMode.CPU_ONLY: cudnn.benchmark = True # Data loading code train_loader, train_sampler, val_loader = create_dataloaders(config) if config.mode.lower() == 'test': print_statistics(compression_algo.statistics()) validate(val_loader, model, criterion, config) if config.mode.lower() == 'train': if not resuming_checkpoint: compression_algo.initialize(train_loader) train(config, compression_algo, model, criterion, is_inception, lr_scheduler, model_name, optimizer, train_loader, train_sampler, val_loader, best_acc1)
def main_worker(current_gpu, config): config.current_gpu = current_gpu config.distributed = config.execution_mode in ( ExecutionMode.DISTRIBUTED, ExecutionMode.MULTIPROCESSING_DISTRIBUTED) if config.distributed: configure_distributed(config) if is_main_process(): configure_logging(config) print_args(config) print(config) config.device = get_device(config) dataset = get_dataset(config.dataset) color_encoding = dataset.color_encoding num_classes = len(color_encoding) weights = config.get('weights') model = load_model(config.model, pretrained=config.get('pretrained', True) if weights is None else False, num_classes=num_classes, model_params=config.get('model_params', {})) compression_algo, model = create_compressed_model(model, config) if weights: sd = torch.load(weights, map_location='cpu') load_state(model, sd) model, model_without_dp = prepare_model_for_execution(model, config) if config.distributed: compression_algo.distributed() resuming_checkpoint = config.resuming_checkpoint if resuming_checkpoint is not None: if not config.pretrained: # Load the previously saved model state model, _, _, _, _ = \ load_checkpoint(model, resuming_checkpoint, config.device, compression_scheduler=compression_algo.scheduler) if config.to_onnx is not None: compression_algo.export_model(config.to_onnx) print("Saved to", config.to_onnx) return if config.mode.lower() == 'test': print(model) model_parameters = filter(lambda p: p.requires_grad, model.parameters()) params = sum([np.prod(p.size()) for p in model_parameters]) print("Trainable argument count:{params}".format(params=params)) model = model.to(config.device) loaders, w_class = load_dataset(dataset, config) _, val_loader = loaders test(model, val_loader, w_class, color_encoding, config) print_statistics(compression_algo.statistics()) elif config.mode.lower() == 'train': loaders, w_class = load_dataset(dataset, config) train_loader, val_loader = loaders if not resuming_checkpoint: compression_algo.initialize(train_loader) model = \ train(model, model_without_dp, compression_algo, train_loader, val_loader, w_class, color_encoding, config) else: # Should never happen...but just in case it does raise RuntimeError( "\"{0}\" is not a valid choice for execution mode.".format( config.mode))
def main_worker(current_gpu, config): ################################# # Setup experiment environment ################################# config.current_gpu = current_gpu config.distributed = config.execution_mode in (ExecutionMode.DISTRIBUTED, ExecutionMode.MULTIPROCESSING_DISTRIBUTED) if config.distributed: configure_distributed(config) if is_on_first_rank(config): configure_logging(config) print_args(config) config.device = get_device(config) config.start_iter = 0 ########################## # Prepare metrics log file ########################## if config.metrics_dump and config.resuming_checkpoint is not None: avg = 0 metrics = {os.path.basename(config.resuming_checkpoint): avg} write_metrics(config, metrics) ################## # Prepare model ################## compression_algo, net = create_model(config) if config.distributed: config.batch_size //= config.ngpus_per_node config.workers //= config.ngpus_per_node compression_algo.distributed() ########################### # Criterion and optimizer ########################### params_to_optimize = get_parameter_groups(net, config) optimizer, lr_scheduler = make_optimizer(params_to_optimize, config) criterion = MultiBoxLoss( config, config['num_classes'], overlap_thresh=0.5, prior_for_matching=True, bkg_label=0, neg_mining=True, neg_pos=3, neg_overlap=0.5, encode_target=False, device=config.device ) ########################### # Load checkpoint ########################### resuming_checkpoint = config.resuming_checkpoint if resuming_checkpoint: print('Resuming training, loading {}...'.format(resuming_checkpoint)) checkpoint = torch.load(resuming_checkpoint, map_location='cpu') # use checkpoint itself in case of only state dict is saved # i.e. checkpoint is created with `torch.save(module.state_dict())` state_dict = checkpoint.get('state_dict', checkpoint) load_state(net, state_dict, is_resume=True) if config.mode.lower() == 'train' and config.to_onnx is None: compression_algo.scheduler.load_state_dict(checkpoint['scheduler']) optimizer.load_state_dict(checkpoint.get('optimizer', optimizer.state_dict())) config.start_iter = checkpoint.get('iter', 0) + 1 if config.to_onnx: compression_algo.export_model(config.to_onnx) print("Saved to {}".format(config.to_onnx)) return ########################### # Prepare data ########################### test_data_loader, train_data_loader = create_dataloaders(config) if config.mode.lower() == 'test': with torch.no_grad(): print_statistics(compression_algo.statistics()) net.eval() mAp = test_net(net, config.device, test_data_loader, distributed=config.distributed) if config.metrics_dump and config.resuming_checkpoint is not None: avg = mAp*100 metrics = {os.path.basename(config.resuming_checkpoint): round(avg, 2)} write_metrics(config, metrics) return if not resuming_checkpoint: compression_algo.initialize(train_data_loader) train(net, compression_algo, train_data_loader, test_data_loader, criterion, optimizer, config, lr_scheduler)