Implement ability to toggle inplace activation operations in the implemented architectures

[javierbg]

🚀 The feature

A lot of the available architectures like AlexNet, EfficientNet, ResNet, DenseNet, etc. are implemented with inplace activation operations, e.g.:

nn.ReLU(inplace=True)

This is done in the name of saving memory (#807)

This could be toggled using a parameter for the model-building functions.

Motivation, pitch

Often, this presents a problem when trying to register register full backward hooks, as you get errors like the following:

UserWarning: Output 0 of BackwardHookFunctionBackward is a view and is being modified inplace. This view was created inside a custom Function (or because an input was returned as-is) and the autograd logic to handle view+inplace would override the custom backward associated with the custom Function, leading to incorrect gradients. This behavior is deprecated and will be forbidden starting version 1.6. You can remove this warning by cloning the output of the custom Function. (Triggered internally at /pytorch/torch/csrc/autograd/variable.cpp:547.)

This is annoying when trying to use libraries like Captum, which need to attach these hooks in layer attribution methods like LRP:

Model cannot contain any in-place nonlinear submodules; these are not supported by the register_full_backward_hook PyTorch API starting from PyTorch v1.9.

Adding a parameter (something like inplace_activations, defaulted to True) would allow to maintain compatibility with old code while providing an option for anyone that needs it.

Alternatives

Currently, my only option is to copy the model code in question (e.g. ResNet) and substituting every inplace=True with inplace=False. This approach works, but it is not really scalable when trying to benchmark as many architectures as possible.

There may be other implementation approaches that libraries like Captum could take in order to be able to implement these methods to work with inplace operations. Nonetheless, the proposed solution appears more simple and would allow a lot of already implemented methods to work out of the box with the standard implementations in torchvision.

Additional context

I'm willing to implement this myself, but I'm having trouble setting up a torchvision development environment in Arch. I might need additional help for that.

cc @datumbox

[datumbox]

@javierbg Thanks for the input. We plan to review how configurable are our model constructors on the future (cc @YosuaMichael), so this is good feedback.

As you correctly said, in-place operators might be a bit faster but can lead to weird results when combined with Eager Quantization or shortcut connections. We have been avoiding adding in-place ops in newer models, except in positions where we feel it's safe to do so. Still, we often need to replace theme with regular relus (during quantization) so we have a private utility method for that:

vision/torchvision/models/quantization/utils.py

Lines 7 to 18 in 8a16e12

	def _replace_relu(module: nn.Module) -> None:
	reassign = {}
	for name, mod in module.named_children():
	_replace_relu(mod)
	# Checking for explicit type instead of instance
	# as we only want to replace modules of the exact type
	# not inherited classes
	if type(mod) is nn.ReLU or type(mod) is nn.ReLU6:
	reassign[name] = nn.ReLU(inplace=False)
	for key, value in reassign.items():
	module._modules[key] = value

So similar replacements are fairly easy to implement for now, within a few lines of code. We will certainly review the model configuration on the future, but until then perhaps adopting a similar approach on your code-base can circumvent the problem? Let me know what you think. Thanks!

[javierbg]

@datumbox Thank you very much for your reply. For sure, what you suggested works flawlessly in my testing, so that's a way better solution than mine. I'm kinda embarrassed that it didn't occur to me before, sometimes I forget about the flexibility of PyTorch!

I'm glad to be of help, so just let me know if you need some work on this once you decide on how to approach model configuration.

[datumbox]

@javierbg No worries at all. PyTorch is massive and has lots of ways to do the same thing. I'll leave this issue open to ensure we review the feedback when we revise the configuration. We'll definitely ping you to provide some help if you are up for it when we kick off the work. In the mean time, have a look on the issues to see if there is anything you would like to contribute.