def prune_finetune_test(iteration, model_cpy, pruning_step, test_fn, train_fn, app_args, param_name, effective_train_size): pylogger = PythonLogger(msglogger) zeros_mask_dict = distiller.create_model_masks_dict(model_cpy) param = model_cpy.state_dict()[param_name] if 0 == prune_tensor(param, param_name, pruning_step, zeros_mask_dict): return (-1, -1, -1, zeros_mask_dict) # Did not prune anything if train_fn is not None: # Fine-tune optimizer = torch.optim.SGD(model_cpy.parameters(), lr=app_args.lr, momentum=app_args.momentum, weight_decay=app_args.weight_decay) app_args.effective_train_size = effective_train_size train_fn(model=model_cpy, compression_scheduler=create_scheduler( model_cpy, zeros_mask_dict), optimizer=optimizer, epoch=iteration, loggers=[pylogger]) # Physically remove filters dataset = app_args.dataset arch = app_args.arch distiller.remove_filters(model_cpy, zeros_mask_dict, arch, dataset, optimizer=None) # Test and record the performance of the pruned model prec1, prec5, loss = test_fn(model=model_cpy, loggers=None) return (prec1, prec5, loss, zeros_mask_dict)
def test_load_gpu_model_on_cpu_with_thinning(): # Issue #148 # 1. create a GPU model and remove 50% of the filters in one of the layers (thninning) # 2. save the thinned model in a checkpoint file # 3. load the checkpoint and place it on the CPU CPU_DEVICE_ID = -1 gpu_model = create_model(False, 'cifar10', 'resnet20_cifar') conv_pname = "module.layer1.0.conv1.weight" conv_p = distiller.model_find_param(gpu_model, conv_pname) pruner = distiller.pruning.L1RankedStructureParameterPruner("test_pruner", group_type="Filters", desired_sparsity=0.5, weights=conv_pname) zeros_mask_dict = distiller.create_model_masks_dict(gpu_model) pruner.set_param_mask(conv_p, conv_pname, zeros_mask_dict, meta=None) # Use the mask to prune zeros_mask_dict[conv_pname].apply_mask(conv_p) distiller.remove_filters(gpu_model, zeros_mask_dict, 'resnet20_cifar', 'cifar10', optimizer=None) assert hasattr(gpu_model, 'thinning_recipes') scheduler = distiller.CompressionScheduler(gpu_model) save_checkpoint(epoch=0, arch='resnet20_cifar', model=gpu_model, scheduler=scheduler, optimizer=None) CPU_DEVICE_ID = -1 cpu_model = create_model(False, 'cifar10', 'resnet20_cifar', device_ids=CPU_DEVICE_ID) load_checkpoint(cpu_model, "checkpoint.pth.tar") assert distiller.model_device(cpu_model) == 'cpu'
def reset(self, init_only=False): """Reset the environment. This is invoked by the Agent. """ msglogger.info("Resetting the environment (init_only={})".format(init_only)) self.current_layer_id = -1 self.prev_action = 0 self.model = copy.deepcopy(self.orig_model) self.zeros_mask_dict = distiller.create_model_masks_dict(self.model) self._remaining_macs = self.dense_model_macs self._removed_macs = 0 if init_only: return obs, _, _, _, = self.step(0) return obs
def reset(self, model): self.model = model self.zeros_mask_dict = distiller.create_model_masks_dict(self.model)
def reset(self, model): self.model = model self.zeros_mask_dict = distiller.create_model_masks_dict(self.model) self.model_metadata = copy.deepcopy(self.cached_model_metadata)