Which function to override if one wants to clip_grad_norm on some, not all, parameters?

[LeeJZh]

❓ Questions and Help

Before asking:

1. Try to find answers to your questions in the Lightning Forum!
2. Search for similar issues.
3. Search the docs.

What is your question?

Which function to override if one wants to clip_grad_norm on some, not all, parameters?

Code

What have you tried?

What's your environment?

• OS: [e.g. iOS, Linux, Win]
• Packaging [e.g. pip, conda]
• Version [e.g. 0.5.2.1]

[tchaton]

Hey there,

Currently in master, the clip_gradients lives in the precision plugin.

from pytorch_lightning.plugins.precision import NativeAMPPlugin

class MyCustomAMPPlugin(NativeAMPPlugin):

    def clip_gradients(self, grad_clip_val: Union[int, float], optimizer: Optimizer, norm_type: float):
        # Example: clip gradients only on the model head
        model = self.trainer.get_model()
        torch.nn.utils.clip_grad_norm_(model.head.parameters(), max_norm=grad_clip_val, norm_type=norm_type)

Trainer(amp_backend='native', precision=16, plugins=[MyCustomAMPPlugin()])

However, this would have to be done for all precision plugin.

We are currently thinking about providing a hook to enable more flexibility there.

@SeanNaren Any idea ?

Best,
T.C

[SeanNaren]

Related to this: #3912

Having clip_gradients as a part of the module makes sense till we realise that different training type/accelerators do different things when clipping gradient norms based on precision. Before this used to lead to us having a bunch of if statements per accelerator in the function within the lightning module, but I think that's not ideal.

I think as you suggested @tchaton is the best way to go, but I might be missing something. In the future I think clip_gradients should live with the training type plugin, but still be overridden as you suggested

[LeeJZh]

Hey there,

Currently in master, the clip_gradients lives in the precision plugin.

from pytorch_lightning.plugins.precision import NativeAMPPlugin

class MyCustomAMPPlugin(NativeAMPPlugin):

    def clip_gradients(self, grad_clip_val: Union[int, float], optimizer: Optimizer, norm_type: float):
        # Example: clip gradients only on the model head
        model = self.trainer.get_model()
        torch.nn.utils.clip_grad_norm_(model.head.parameters(), max_norm=grad_clip_val, norm_type=norm_type)

Trainer(amp_backend='native', precision=16, plugins=[MyCustomAMPPlugin()])

However, this would have to be done for all precision plugin.

We are currently thinking about providing a hook to enable more flexibility there.

@SeanNaren Any idea ?

Best,
T.C

Thanks, that helps a lot!

[LeeJZh]

Hey there,

Currently in master, the clip_gradients lives in the precision plugin.

from pytorch_lightning.plugins.precision import NativeAMPPlugin

class MyCustomAMPPlugin(NativeAMPPlugin):

    def clip_gradients(self, grad_clip_val: Union[int, float], optimizer: Optimizer, norm_type: float):
        # Example: clip gradients only on the model head
        model = self.trainer.get_model()
        torch.nn.utils.clip_grad_norm_(model.head.parameters(), max_norm=grad_clip_val, norm_type=norm_type)

Trainer(amp_backend='native', precision=16, plugins=[MyCustomAMPPlugin()])

However, this would have to be done for all precision plugin.

We are currently thinking about providing a hook to enable more flexibility there.

@SeanNaren Any idea ?

Best,
T.C

Hi, I keep getting

'NoneType' object has no attribute 'scaler'

error

with

from pytorch_lightning.plugins.native_amp import NativeAMPPlugin

class MyCustomAMPPlugin(NativeAMPPlugin):

    def clip_gradients(self, grad_clip_val: Union[int, float], optimizer: Optimizer, norm_type: float):

        model = self.trainer.get_model()
        curr_epoch = self.trainer.current_epoch
        if curr_epoch < 30:
            gcv = grad_clip_val * 2
        elif curr_epoch < 60:
            gcv = grad_clip_val * 1.5
        else:
            gcv = grad_clip_val
        torch.nn.utils.clip_grad_norm_(model.scout_parameters, max_norm=gcv, norm_type=norm_type)

trainer = pl.Trainer.from_argparse_args(args,
                                            logger=logger,
                                            callbacks=[checkpoint_callback], plugins=[MyCustomAMPPlugin()])

i guess i fail to pass tranier to plugins?

[stale]

This issue has been automatically marked as stale because it hasn't had any recent activity. This issue will be closed in 7 days if no further activity occurs. Thank you for your contributions, Pytorch Lightning Team!