Problems in Pruning

[joshclancy]

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?

[carmocca]

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:

https://github.com/PyTorchLightning/pytorch-lightning/blob/4c07ab5e99dd20c1f309d9e73cdaacc1ebad9499/pytorch_lightning/callbacks/pruning.py#L422-L444

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.

[joshclancy]

Thanks got the reply. pl_module.modules returns <generator object Module.modules at 0x7f5ca44f3dd0> Does this mean anything to you?

I'll go ahead and try switching to master branch.

[carmocca]

Sorry, I should've said list(pl_module.modules())

[joshclancy]

Ahh, that returns:

[LitModel(
  (model): AutoEncoder_sigma(
    (conv1): Conv2d(3, 32, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
    (conv2): Conv2d(32, 64, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
    (lin1): Linear(in_features=4096, out_features=128, bias=True)
    (lin2): Linear(in_features=128, out_features=4096, bias=True)
    (upconv1): ConvTranspose2d(64, 32, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
    (upconv2): ConvTranspose2d(32, 3, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
  )
), AutoEncoder_sigma(
  (conv1): Conv2d(3, 32, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
  (conv2): Conv2d(32, 64, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
  (lin1): Linear(in_features=4096, out_features=128, bias=True)
  (lin2): Linear(in_features=128, out_features=4096, bias=True)
  (upconv1): ConvTranspose2d(64, 32, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
  (upconv2): ConvTranspose2d(32, 3, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
), Conv2d(3, 32, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1)), Conv2d(32, 64, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1)), Linear(in_features=4096, out_features=128, bias=True), Linear(in_features=128, out_features=4096, bias=True), ConvTranspose2d(64, 32, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1)), ConvTranspose2d(32, 3, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))]