Problems in Pruning
Hello, I am trying to get pruning to work within my lightning model. I have tried multiple methods, but have not been able to get the ModelPruning module to work.
When I try this:
def compute_amount(epoch):
    # the sum of all returned values need to be smaller than 1
    if epoch == 10:
        return 0.5

    elif epoch == 50:
        return 0.25

    elif 75 < epoch < 99 :
        return 0.01
    else:
      return 0

prune = ModelPruning(
            pruning_fn='l1_unstructured',
            parameters_to_prune=[(model.model.conv1, 'weight'),
                            (model.model.conv2, 'weight'),
                            (model.model.lin1, 'weight'),
                            (model.model.lin2, 'weight'),
                            (model.model.upconv1, 'weight'),
                            (model.model.upconv2, 'weight')],
            amount=compute_amount,
            use_global_unstructured=True,
        )

trainer = pl.Trainer(gpus=-1, max_epochs = 15, logger=wandb_logger, profiler=SimpleProfiler("/content/bla"), progress_bar_refresh_rate=20, callbacks=[prune])
trainer.fit(model, dm)

I get this output:

MisconfigurationException                 Traceback (most recent call last)
 in ()
11
12 trainer = pl.Trainer(gpus=-1, max_epochs = 15, logger=wandb_logger, profiler=SimpleProfiler("/content/bla"), progress_bar_refresh_rate=20, callbacks=[prune])
---> 13 trainer.fit(model, dm)
4 frames
/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/trainer.py in fit(self, model, train_dataloader, val_dataloaders, datamodule)
457         # ----------------------------
458         self.call_setup_hook(model)
--> 459         self.call_hook("on_before_accelerator_backend_setup", model)
460         self.accelerator.setup(self, model)  # note: this sets up self.lightning_module
461
/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/trainer.py in call_hook(self, hook_name, *args, **kwargs)
1089             if hasattr(self, hook_name):
1090                 trainer_hook = getattr(self, hook_name)
-> 1091                 trainer_hook(*args, **kwargs)
1092
1093             # next call hook in lightningModule
/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/callback_hook.py in on_before_accelerator_backend_setup(self, model)
33         """Called in the beginning of fit and test"""
34         for callback in self.callbacks:
---> 35             callback.on_before_accelerator_backend_setup(self, model)
36
37     def setup(self, model, stage: str):
/usr/local/lib/python3.7/dist-packages/pytorch_lightning/callbacks/pruning.py in on_before_accelerator_backend_setup(self, trainer, pl_module)
360     def on_before_accelerator_backend_setup(self, trainer, pl_module: LightningModule):
361         parameters_to_prune = self.sanitize_parameters_to_prune(
--> 362             pl_module, self._parameters_to_prune, parameter_names=self._parameter_names
363         )
364
/usr/local/lib/python3.7/dist-packages/pytorch_lightning/callbacks/pruning.py in sanitize_parameters_to_prune(pl_module, parameters_to_prune, parameter_names)
441             if missing_modules or missing_parameters:
442                 raise MisconfigurationException(
--> 443                     "Some provided parameters_to_tune don't exist in the model."
444                     f" Found missing modules: {missing_modules} and missing parameters: {missing_parameters}"
445                 )
MisconfigurationException: Some provided parameters_to_tune don't exist in the model. Found missing modules: [Conv2d(1, 8, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1)), Conv2d(8, 16, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1)), Linear(in_features=784, out_features=32, bias=True), Linear(in_features=32, out_features=784, bias=True), ConvTranspose2d(16, 8, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1)), ConvTranspose2d(8, 1, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))] and missing parameters: []

Which makes it seem like ModelPruning is not finding my model's weights, though I do not understand why. If I put into ModelPruning weights that do not exist in the model. It catches that problem during ModelPruning's init. Which is not happening here. So that does not seem to be the problem.
I have also tried bypassing the parameters_to_prune parameter altogether.
def compute_amount(epoch):
    # the sum of all returned values need to be smaller than 1
    if epoch == 10:
        return 0.5

    elif epoch == 50:
        return 0.25

    elif 75 < epoch < 99 :
        return 0.01
    else:
      return 0

prune = ModelPruning(
            pruning_fn='l1_unstructured',
            parameter_names=['weight', 'bias'],
            amount=compute_amount,
            use_global_unstructured=True,
        )

trainer = pl.Trainer(gpus=-1, max_epochs = 15, logger=wandb_logger, profiler=SimpleProfiler("/content/bla"), progress_bar_refresh_rate=20, callbacks=[prune])
trainer.fit(model, dm)

This works right until it does not. One epoch runs and then it snags on a callback within the validation step.


TypeError                                 Traceback (most recent call last)
/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/trainer.py in run_train(self)
636                     # run train epoch
--> 637                     self.train_loop.run_training_epoch()
638
56 frames
/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/training_loop.py in run_training_epoch(self)
576         if should_check_val:
--> 577             self.trainer.run_evaluation(on_epoch=True)
578
/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/trainer.py in run_evaluation(self, max_batches, on_epoch)
750         # hook
--> 751         self.evaluation_loop.on_evaluation_end()
752
/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/evaluation_loop.py in on_evaluation_end(self, *args, **kwargs)
99         else:
--> 100             self.trainer.call_hook('on_validation_end', *args, **kwargs)
101
/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/trainer.py in call_hook(self, hook_name, *args, **kwargs)
1090                 trainer_hook = getattr(self, hook_name)
-> 1091                 trainer_hook(*args, **kwargs)
1092
/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/callback_hook.py in on_validation_end(self)
184         for callback in self.callbacks:
--> 185             callback.on_validation_end(self, self.lightning_module)
186
/usr/local/lib/python3.7/dist-packages/pytorch_lightning/callbacks/model_checkpoint.py in on_validation_end(self, trainer, pl_module)
211         """
--> 212         self.save_checkpoint(trainer, pl_module)
213
/usr/local/lib/python3.7/dist-packages/pytorch_lightning/callbacks/model_checkpoint.py in save_checkpoint(self, trainer, pl_module)
261         # Mode 2: save the last checkpoint
--> 262         self._save_last_checkpoint(trainer, pl_module, monitor_candidates)
263
/usr/local/lib/python3.7/dist-packages/pytorch_lightning/callbacks/model_checkpoint.py in _save_last_checkpoint(self, trainer, pl_module, ckpt_name_metrics)
545         else:
--> 546             self._save_model(last_filepath, trainer, pl_module)
547         if (
/usr/local/lib/python3.7/dist-packages/pytorch_lightning/callbacks/model_checkpoint.py in _save_model(self, filepath, trainer, pl_module)
334         if self.save_function is not None:
--> 335             self.save_function(filepath, self.save_weights_only)
336         else:
/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/properties.py in save_checkpoint(self, filepath, weights_only)
326     def save_checkpoint(self, filepath, weights_only: bool = False) -> None:
--> 327         self.checkpoint_connector.save_checkpoint(filepath, weights_only)
328
/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/connectors/checkpoint_connector.py in save_checkpoint(self, filepath, weights_only)
396         # dump states as a checkpoint dictionary object
--> 397         checkpoint = self.dump_checkpoint(weights_only)
398         if self.trainer.is_global_zero:
/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/connectors/checkpoint_connector.py in dump_checkpoint(self, weights_only)
283             # dump callbacks
--> 284             checkpoint['callbacks'] = self.trainer.on_save_checkpoint(checkpoint)
285
/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/callback_hook.py in on_save_checkpoint(self, checkpoint)
221             else:
--> 222                 state = callback.on_save_checkpoint(self, self.lightning_module, checkpoint)
223             if state:
/usr/local/lib/python3.7/dist-packages/pytorch_lightning/callbacks/pruning.py in on_save_checkpoint(self, trainer, pl_module, checkpoint)
399             # prune a copy so training can continue with the same buffers
--> 400             copy = deepcopy(pl_module.to("cpu"))
401             self.make_pruning_permanent(copy)
/usr/lib/python3.7/copy.py in deepcopy(x, memo, _nil)
179                 else:
--> 180                     y = _reconstruct(x, memo, *rv)
181
/usr/lib/python3.7/copy.py in _reconstruct(x, memo, func, args, state, listiter, dictiter, deepcopy)
280         if deep:
--> 281             state = deepcopy(state, memo)
282         if hasattr(y, 'setstate'):
/usr/lib/python3.7/copy.py in deepcopy(x, memo, _nil)
149     if copier:
--> 150         y = copier(x, memo)
151     else:
/usr/lib/python3.7/copy.py in _deepcopy_dict(x, memo, deepcopy)
240     for key, value in x.items():
--> 241         y[deepcopy(key, memo)] = deepcopy(value, memo)
242     return y
/usr/lib/python3.7/copy.py in deepcopy(x, memo, _nil)
179                 else:
--> 180                     y = _reconstruct(x, memo, *rv)
181
/usr/lib/python3.7/copy.py in _reconstruct(x, memo, func, args, state, listiter, dictiter, deepcopy)
280         if deep:
--> 281             state = deepcopy(state, memo)
282         if hasattr(y, 'setstate'):
/usr/lib/python3.7/copy.py in deepcopy(x, memo, _nil)
149     if copier:
--> 150         y = copier(x, memo)
151     else:
/usr/lib/python3.7/copy.py in _deepcopy_dict(x, memo, deepcopy)
240     for key, value in x.items():
--> 241         y[deepcopy(key, memo)] = deepcopy(value, memo)
242     return y
/usr/lib/python3.7/copy.py in deepcopy(x, memo, _nil)
179                 else:
--> 180                     y = _reconstruct(x, memo, *rv)
181
/usr/lib/python3.7/copy.py in _reconstruct(x, memo, func, args, state, listiter, dictiter, deepcopy)
280         if deep:
--> 281             state = deepcopy(state, memo)
282         if hasattr(y, 'setstate'):
/usr/lib/python3.7/copy.py in deepcopy(x, memo, _nil)
149     if copier:
--> 150         y = copier(x, memo)
151     else:
/usr/lib/python3.7/copy.py in _deepcopy_dict(x, memo, deepcopy)
240     for key, value in x.items():
--> 241         y[deepcopy(key, memo)] = deepcopy(value, memo)
242     return y
/usr/lib/python3.7/copy.py in deepcopy(x, memo, _nil)
168                     if reductor:
--> 169                         rv = reductor(4)
170                     else:
TypeError: cannot serialize '_io.TextIOWrapper' object
During handling of the above exception, another exception occurred:
TypeError                                 Traceback (most recent call last)
 in ()
11
12 trainer = pl.Trainer(gpus=-1, max_epochs = 15, logger=wandb_logger, profiler=SimpleProfiler("/content/bla"), progress_bar_refresh_rate=20, callbacks=[prune])
---> 13 trainer.fit(model, dm)
/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/trainer.py in fit(self, model, train_dataloader, val_dataloaders, datamodule)
497
498         # dispath start_training or start_testing or start_predicting
--> 499         self.dispatch()
500
501         # plugin will finalized fitting (e.g. ddp_spawn will load trained model)
/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/trainer.py in dispatch(self)
544
545         else:
--> 546             self.accelerator.start_training(self)
547
548     def train_or_test_or_predict(self):
/usr/local/lib/python3.7/dist-packages/pytorch_lightning/accelerators/accelerator.py in start_training(self, trainer)
71
72     def start_training(self, trainer):
---> 73         self.training_type_plugin.start_training(trainer)
74
75     def start_testing(self, trainer):
/usr/local/lib/python3.7/dist-packages/pytorch_lightning/plugins/training_type/training_type_plugin.py in start_training(self, trainer)
112     def start_training(self, trainer: 'Trainer') -> None:
113         # double dispatch to initiate the training loop
--> 114         self._results = trainer.run_train()
115
116     def start_testing(self, trainer: 'Trainer') -> None:
/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/trainer.py in run_train(self)
667         finally:
668             # hook
--> 669             self.train_loop.on_train_end()
670
671     def run_evaluation(self, max_batches=None, on_epoch=False):
/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/training_loop.py in on_train_end(self)
132         # when a checkpoint was saved at the last step
133         self.trainer.global_step -= 1
--> 134         self.check_checkpoint_callback(should_update=True, is_last=True)
135         self.trainer.global_step += 1
136
/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/training_loop.py in check_checkpoint_callback(self, should_update, is_last)
162
163             for cb in callbacks:
--> 164                 cb.on_validation_end(self.trainer, model)
165
166     def check_early_stopping_callback(self, should_update):
/usr/local/lib/python3.7/dist-packages/pytorch_lightning/callbacks/model_checkpoint.py in on_validation_end(self, trainer, pl_module)
210         checkpoints can be saved at the end of the val loop
211         """
--> 212         self.save_checkpoint(trainer, pl_module)
213
214     def on_save_checkpoint(self, trainer, pl_module, checkpoint: Dict[str, Any]) -> Dict[str, Any]:
/usr/local/lib/python3.7/dist-packages/pytorch_lightning/callbacks/model_checkpoint.py in save_checkpoint(self, trainer, pl_module)
260
261         # Mode 2: save the last checkpoint
--> 262         self._save_last_checkpoint(trainer, pl_module, monitor_candidates)
263
264     def __validate_init_configuration(self):
/usr/local/lib/python3.7/dist-packages/pytorch_lightning/callbacks/model_checkpoint.py in _save_last_checkpoint(self, trainer, pl_module, ckpt_name_metrics)
544             trainer.training_type_plugin.rpc_save_model(self._save_model, last_filepath, trainer, pl_module)
545         else:
--> 546             self._save_model(last_filepath, trainer, pl_module)
547         if (
548             self.last_model_path and self.last_model_path != last_filepath
/usr/local/lib/python3.7/dist-packages/pytorch_lightning/callbacks/model_checkpoint.py in _save_model(self, filepath, trainer, pl_module)
333         # delegate the saving to the trainer
334         if self.save_function is not None:
--> 335             self.save_function(filepath, self.save_weights_only)
336         else:
337             raise ValueError(".save_function() not set")
/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/properties.py in save_checkpoint(self, filepath, weights_only)
325
326     def save_checkpoint(self, filepath, weights_only: bool = False) -> None:
--> 327         self.checkpoint_connector.save_checkpoint(filepath, weights_only)
328
329     @Property
/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/connectors/checkpoint_connector.py in save_checkpoint(self, filepath, weights_only)
395         """
396         # dump states as a checkpoint dictionary object
--> 397         checkpoint = self.dump_checkpoint(weights_only)
398         if self.trainer.is_global_zero:
399             # write the checkpoint dictionary on the file
/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/connectors/checkpoint_connector.py in dump_checkpoint(self, weights_only)
282         if not weights_only:
283             # dump callbacks
--> 284             checkpoint['callbacks'] = self.trainer.on_save_checkpoint(checkpoint)
285
286             optimizer_states = []
/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/callback_hook.py in on_save_checkpoint(self, checkpoint)
220                 state = callback.on_save_checkpoint(self, self.lightning_module)  # noqa: parameter-unfilled
221             else:
--> 222                 state = callback.on_save_checkpoint(self, self.lightning_module, checkpoint)
223             if state:
224                 callback_states[type(callback)] = state
/usr/local/lib/python3.7/dist-packages/pytorch_lightning/callbacks/pruning.py in on_save_checkpoint(self, trainer, pl_module, checkpoint)
398             prev_device = pl_module.device
399             # prune a copy so training can continue with the same buffers
--> 400             copy = deepcopy(pl_module.to("cpu"))
401             self.make_pruning_permanent(copy)
402             checkpoint["state_dict"] = copy.state_dict()
/usr/lib/python3.7/copy.py in deepcopy(x, memo, _nil)
178                     y = x
179                 else:
--> 180                     y = _reconstruct(x, memo, *rv)
181
182     # If is its own copy, don't memoize.
/usr/lib/python3.7/copy.py in _reconstruct(x, memo, func, args, state, listiter, dictiter, deepcopy)
279     if state is not None:
280         if deep:
--> 281             state = deepcopy(state, memo)
282         if hasattr(y, 'setstate'):
283             y.setstate(state)
/usr/lib/python3.7/copy.py in deepcopy(x, memo, _nil)
148     copier = _deepcopy_dispatch.get(cls)
149     if copier:
--> 150         y = copier(x, memo)
151     else:
152         try:
/usr/lib/python3.7/copy.py in _deepcopy_dict(x, memo, deepcopy)
239     memo[id(x)] = y
240     for key, value in x.items():
--> 241         y[deepcopy(key, memo)] = deepcopy(value, memo)
242     return y
243 d[dict] = _deepcopy_dict
/usr/lib/python3.7/copy.py in deepcopy(x, memo, _nil)
178                     y = x
179                 else:
--> 180                     y = _reconstruct(x, memo, *rv)
181
182     # If is its own copy, don't memoize.
/usr/lib/python3.7/copy.py in _reconstruct(x, memo, func, args, state, listiter, dictiter, deepcopy)
279     if state is not None:
280         if deep:
--> 281             state = deepcopy(state, memo)
282         if hasattr(y, 'setstate'):
283             y.setstate(state)
/usr/lib/python3.7/copy.py in deepcopy(x, memo, _nil)
148     copier = _deepcopy_dispatch.get(cls)
149     if copier:
--> 150         y = copier(x, memo)
151     else:
152         try:
/usr/lib/python3.7/copy.py in _deepcopy_dict(x, memo, deepcopy)
239     memo[id(x)] = y
240     for key, value in x.items():
--> 241         y[deepcopy(key, memo)] = deepcopy(value, memo)
242     return y
243 d[dict] = _deepcopy_dict
/usr/lib/python3.7/copy.py in deepcopy(x, memo, _nil)
178                     y = x
179                 else:
--> 180                     y = _reconstruct(x, memo, *rv)
181
182     # If is its own copy, don't memoize.
/usr/lib/python3.7/copy.py in _reconstruct(x, memo, func, args, state, listiter, dictiter, deepcopy)
279     if state is not None:
280         if deep:
--> 281             state = deepcopy(state, memo)
282         if hasattr(y, 'setstate'):
283             y.setstate(state)
/usr/lib/python3.7/copy.py in deepcopy(x, memo, _nil)
148     copier = _deepcopy_dispatch.get(cls)
149     if copier:
--> 150         y = copier(x, memo)
151     else:
152         try:
/usr/lib/python3.7/copy.py in _deepcopy_dict(x, memo, deepcopy)
239     memo[id(x)] = y
240     for key, value in x.items():
--> 241         y[deepcopy(key, memo)] = deepcopy(value, memo)
242     return y
243 d[dict] = _deepcopy_dict
/usr/lib/python3.7/copy.py in deepcopy(x, memo, _nil)
167                     reductor = getattr(x, "reduce_ex", None)
168                     if reductor:
--> 169                         rv = reductor(4)
170                     else:
171                         reductor = getattr(x, "reduce", None)
TypeError: cannot serialize '_io.TextIOWrapper' object

I am a bit overwhelmed by that error. Does anyone know what I am doing wrong? Or have I found a bug?
For the first error, what does pl_module.modules() return?
You get the error because the following code can't find the parameters you passed in your model:

  
    
      pytorch-lightning/pytorch_lightning/callbacks/pruning.py
    
    
        Lines 422 to 444
      in
      4c07ab5
    
  
  
    

        
          
           current_modules = [m for m in pl_module.modules() if not isinstance(m, _MODULE_CONTAINERS)] 
        

        
          
            
        

        
          
           if parameters_to_prune is None: 
        

        
          
               parameters_to_prune = [(m, p) for p in parameters for m in current_modules 
        

        
          
                                      if getattr(m, p, None) is not None] 
        

        
          
           elif ( 
        

        
          
               isinstance(parameters_to_prune, (list, tuple)) and len(parameters_to_prune) > 0 
        

        
          
               and all(len(p) == 2 for p in parameters_to_prune) 
        

        
          
               and all(isinstance(a, nn.Module) and isinstance(b, str) for a, b in parameters_to_prune) 
        

        
          
           ): 
        

        
          
               missing_modules, missing_parameters = [], [] 
        

        
          
               for module, name in parameters_to_prune: 
        

        
          
                   if module not in current_modules: 
        

        
          
                       missing_modules.append(module) 
        

        
          
                       continue 
        

        
          
                   if not hasattr(module, name): 
        

        
          
                       missing_parameters.append(name) 
        

        
          
            
        

        
          
               if missing_modules or missing_parameters: 
        

        
          
                   raise MisconfigurationException( 
        

        
          
                       "Some provided `parameters_to_tune` don't exist in the model." 
        

        
          
                       f" Found missing modules: {missing_modules} and missing parameters: {missing_parameters}" 
        

        
          
                   ) 
        
    
  


The second error is definitely a bug. It happens because you are also using a profiler. Should be fixed in master, the fix will be released in 1.3. Feel free to try it on master in the meantime.