Spectral normalization in Opacus

[100sunnygirl]

Hello! I have a question about using spectral norm with this framework.

I want to use torch.nn.utils.spectral_norm for my GAN model to stabilize training. But when I use this function to the layers(including nn.Linear and nn.Conv2d) in the GAN model, the GradSampleModule can't support all the layers. So how can I use spectral norm with Opacus?

Expecting your response!

[HuanyuZhang]

Could you link the error message here? thx!

[leesocool]

Hey, I'm having this issue too. When I use torch.nn.utils.spectral_norm, the error occurs as follows :

NotImplementedError: [NotImplementedError("Model contains a trainable layer that Opacus doesn't currently support(blocks.0.0.conv2d0.conv2d:Conv2d(3, 192, kernel_size=(1, 1), stride=(1, 1))). Please implement and register grad sampler for this layer. (See opacus.grad_sample.utils.register_grad_sampler)"), NotImplementedError("Model contains a trainable layer that Opacus doesn't currently support(blocks.0.0.conv2d1.conv2d:Conv2d(3, 192, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))). Please implement and register grad sampler for this layer. (See opacus.grad_sample.utils.register_grad_sampler)"), NotImplementedError("Model contains a trainable layer that Opacus doesn't currently support(blocks.0.0.conv2d2.conv2d:Conv2d(192, 192, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))). Please implement and register grad sampler for this layer. (See opacus.grad_sample.utils.register_grad_sampler)"), NotImplementedError("Model contains a trainable layer that Opacus doesn't currently support(blocks.1.0:SelfAttention(\n (conv1x1_theta): SNConv2d(\n (conv2d): Conv2d(192, 24, kernel_size=(1, 1), stride=(1, 1), bias=False)\n )\n (conv1x1_phi): SNConv2d(\n (conv2d): Conv2d(192, 24, kernel_size=(1, 1), stride=(1, 1), bias=False)\n )\n (conv1x1_g): SNConv2d(\n (conv2d): Conv2d(192, 96, kernel_size=(1, 1), stride=(1, 1), bias=False)\n )\n (conv1x1_attn): SNConv2d(\n (conv2d): Conv2d(96, 192, kernel_size=(1, 1), stride=(1, 1), bias=False)\n )\n (maxpool): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n (softmax): Softmax(dim=-1)\n)). Please implement and register grad sampler for this layer. (See opacus.grad_sample.utils.register_grad_sampler)"), NotImplementedError("Model contains a trainable layer that Opacus doesn't currently support(blocks.1.0.conv1x1_theta.conv2d:Conv2d(192, 24, kernel_size=(1, 1), stride=(1, 1), bias=False)). Please implement and register grad sampler for this layer. (See opacus.grad_sample.utils.register_grad_sampler)"), NotImplementedError("Model contains a trainable layer that Opacus doesn't currently support(blocks.1.0.conv1x1_phi.conv2d:Conv2d(192, 24, kernel_size=(1, 1), stride=(1, 1), bias=False)). Please implement and register grad sampler for this layer. (See opacus.grad_sample.utils.register_grad_sampler)"), NotImplementedError("Model contains a trainable layer that Opacus doesn't currently support(blocks.1.0.conv1x1_g.conv2d:Conv2d(192, 96, kernel_size=(1, 1), stride=(1, 1), bias=False)). Please implement and register grad sampler for this layer. (See opacus.grad_sample.utils.register_grad_sampler)"), NotImplementedError("Model contains a trainable layer that Opacus doesn't currently support(blocks.1.0.conv1x1_attn.conv2d:Conv2d(96, 192, kernel_size=(1, 1), stride=(1, 1), bias=False)). Please implement and register grad sampler for this layer. (See opacus.grad_sample.utils.register_grad_sampler)"), NotImplementedError("Model contains a trainable layer that Opacus doesn't currently support(blocks.2.0.conv2d0.conv2d:Conv2d(192, 192, kernel_size=(1, 1), stride=(1, 1))). Please implement and register grad sampler for this layer. (See opacus.grad_sample.utils.register_grad_sampler)"), NotImplementedError("Model contains a trainable layer that Opacus doesn't currently support(blocks.2.0.conv2d1.conv2d:Conv2d(192, 192, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))). Please implement and register grad sampler for this layer. (See opacus.grad_sample.utils.register_grad_sampler)"), NotImplementedError("Model contains a trainable layer that Opacus doesn't currently support(blocks.2.0.conv2d2.conv2d:Conv2d(192, 192, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))). Please implement and register grad sampler for this layer. (See opacus.grad_sample.utils.register_grad_sampler)"), NotImplementedError("Model contains a trainable layer that Opacus doesn't currently support(blocks.3.0.conv2d1.conv2d:Conv2d(192, 192, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))). Please implement and register grad sampler for this layer. (See opacus.grad_sample.utils.register_grad_sampler)"), NotImplementedError("Model contains a trainable layer that Opacus doesn't currently support(blocks.3.0.conv2d2.conv2d:Conv2d(192, 192, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))). Please implement and register grad sampler for this layer. (See opacus.grad_sample.utils.register_grad_sampler)"), NotImplementedError("Model contains a trainable layer that Opacus doesn't currently support(blocks.4.0.conv2d1.conv2d:Conv2d(192, 192, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))). Please implement and register grad sampler for this layer. (See opacus.grad_sample.utils.register_grad_sampler)"), NotImplementedError("Model contains a trainable layer that Opacus doesn't currently support(blocks.4.0.conv2d2.conv2d:Conv2d(192, 192, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))). Please implement and register grad sampler for this layer. (See opacus.grad_sample.utils.register_grad_sampler)"), NotImplementedError("Model contains a trainable layer that Opacus doesn't currently support(linear1.liner:Linear(in_features=192, out_features=1, bias=True)). Please implement and register grad sampler for this layer. (See opacus.grad_sample.utils.register_grad_sampler)")]

When I use torch.nn.utils.parametrizations.spectral_norm, the error occurs as follows :

NotImplementedError: [NotImplementedError("Model contains a trainable layer that Opacus doesn't currently support(blocks.0.0.conv2d0.conv2d:ParametrizedConv2d(\n 3, 192, kernel_size=(1, 1), stride=(1, 1)\n (parametrizations): ModuleDict(\n (weight): ParametrizationList(\n (0): _SpectralNorm()\n )\n )\n)). Please implement and register grad sampler for this layer. (See opacus.grad_sample.utils.register_grad_sampler)"), NotImplementedError("Model contains a trainable layer that Opacus doesn't currently support(blocks.0.0.conv2d0.conv2d.parametrizations.weight:ParametrizationList(\n (0): _SpectralNorm()\n)). Please implement and register grad sampler for this layer. (See opacus.grad_sample.utils.register_grad_sampler)"), NotImplementedError("Model contains a trainable layer that Opacus doesn't currently support(blocks.0.0.conv2d1.conv2d:ParametrizedConv2d(\n 3, 192, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)\n (parametrizations): ModuleDict(\n (weight): ParametrizationList(\n (0): _SpectralNorm()\n )\n )\n)). Please implement and register grad sampler for this layer. (See opacus.grad_sample.utils.register_grad_sampler)"), NotImplementedError("Model contains a trainable layer that Opacus doesn't currently support(blocks.0.0.conv2d1.conv2d.parametrizations.weight:ParametrizationList(\n (0): _SpectralNorm()\n)). Please implement and register grad sampler for this layer. (See opacus.grad_sample.utils.register_grad_sampler)"), NotImplementedError("Model contains a trainable layer that Opacus doesn't currently support(blocks.0.0.conv2d2.conv2d:ParametrizedConv2d(\n 192, 192, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)\n (parametrizations): ModuleDict(\n (weight): ParametrizationList(\n (0): _SpectralNorm()\n )\n )\n)). Please implement and register grad sampler for this layer. (See opacus.grad_sample.utils.register_grad_sampler)"), NotImplementedError("Model contains a trainable layer that Opacus doesn't currently support(blocks.0.0.conv2d2.conv2d.parametrizations.weight:ParametrizationList(\n (0): _SpectralNorm()\n)). Please implement and register grad sampler for this layer. (See opacus.grad_sample.utils.register_grad_sampler)"), NotImplementedError("Model contains a trainable layer that Opacus doesn't currently support(blocks.1.0:SelfAttention(\n (conv1x1_theta): SNConv2d(\n (conv2d): ParametrizedConv2d(\n 192, 24, kernel_size=(1, 1), stride=(1, 1), bias=False\n (parametrizations): ModuleDict(\n (weight): ParametrizationList(\n (0): _SpectralNorm()\n )\n )\n )\n )\n (conv1x1_phi): SNConv2d(\n (conv2d): ParametrizedConv2d(\n 192, 24, kernel_size=(1, 1), stride=(1, 1), bias=False\n (parametrizations): ModuleDict(\n (weight): ParametrizationList(\n (0): _SpectralNorm()\n )\n )\n )\n )\n (conv1x1_g): SNConv2d(\n (conv2d): ParametrizedConv2d(\n 192, 96, kernel_size=(1, 1), stride=(1, 1), bias=False\n (parametrizations): ModuleDict(\n (weight): ParametrizationList(\n (0): _SpectralNorm()\n )\n )\n )\n )\n (conv1x1_attn): SNConv2d(\n (conv2d): ParametrizedConv2d(\n 96, 192, kernel_size=(1, 1), stride=(1, 1), bias=False\n (parametrizations): ModuleDict(\n (weight): ParametrizationList(\n (0): _SpectralNorm()\n )\n )\n )\n )\n (maxpool): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n (softmax): Softmax(dim=-1)\n)). Please implement and register grad sampler for this layer. (See opacus.grad_sample.utils.register_grad_sampler)"), NotImplementedError("Model contains a trainable layer that Opacus doesn't currently support(blocks.1.0.conv1x1_theta.conv2d.parametrizations.weight:ParametrizationList(\n (0): _SpectralNorm()\n)). Please implement and register grad sampler for this layer. (See opacus.grad_sample.utils.register_grad_sampler)"), NotImplementedError("Model contains a trainable layer that Opacus doesn't currently support(blocks.1.0.conv1x1_phi.conv2d.parametrizations.weight:ParametrizationList(\n (0): _SpectralNorm()\n)). Please implement and register grad sampler for this layer. (See opacus.grad_sample.utils.register_grad_sampler)"), NotImplementedError("Model contains a trainable layer that Opacus doesn't currently support(blocks.1.0.conv1x1_g.conv2d.parametrizations.weight:ParametrizationList(\n (0): _SpectralNorm()\n)). Please implement and register grad sampler for this layer. (See opacus.grad_sample.utils.register_grad_sampler)"), NotImplementedError("Model contains a trainable layer that Opacus doesn't currently support(blocks.1.0.conv1x1_attn.conv2d.parametrizations.weight:ParametrizationList(\n (0): _SpectralNorm()\n)). Please implement and register grad sampler for this layer. (See opacus.grad_sample.utils.register_grad_sampler)"), NotImplementedError("Model contains a trainable layer that Opacus doesn't currently support(blocks.2.0.conv2d0.conv2d:ParametrizedConv2d(\n 192, 192, kernel_size=(1, 1), stride=(1, 1)\n (parametrizations): ModuleDict(\n (weight): ParametrizationList(\n (0): _SpectralNorm()\n )\n )\n)). Please implement and register grad sampler for this layer. (See opacus.grad_sample.utils.register_grad_sampler)"), NotImplementedError("Model contains a trainable layer that Opacus doesn't currently support(blocks.2.0.conv2d0.conv2d.parametrizations.weight:ParametrizationList(\n (0): _SpectralNorm()\n)). Please implement and register grad sampler for this layer. (See opacus.grad_sample.utils.register_grad_sampler)"), NotImplementedError("Model contains a trainable layer that Opacus doesn't currently support(blocks.2.0.conv2d1.conv2d:ParametrizedConv2d(\n 192, 192, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)\n (parametrizations): ModuleDict(\n (weight): ParametrizationList(\n (0): _SpectralNorm()\n )\n )\n)). Please implement and register grad sampler for this layer. (See opacus.grad_sample.utils.register_grad_sampler)"), NotImplementedError("Model contains a trainable layer that Opacus doesn't currently support(blocks.2.0.conv2d1.conv2d.parametrizations.weight:ParametrizationList(\n (0): _SpectralNorm()\n)). Please implement and register grad sampler for this layer. (See opacus.grad_sample.utils.register_grad_sampler)"), NotImplementedError("Model contains a trainable layer that Opacus doesn't currently support(blocks.2.0.conv2d2.conv2d:ParametrizedConv2d(\n 192, 192, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)\n (parametrizations): ModuleDict(\n (weight): ParametrizationList(\n (0): _SpectralNorm()\n )\n )\n)). Please implement and register grad sampler for this layer. (See opacus.grad_sample.utils.register_grad_sampler)"), NotImplementedError("Model contains a trainable layer that Opacus doesn't currently support(blocks.2.0.conv2d2.conv2d.parametrizations.weight:ParametrizationList(\n (0): _SpectralNorm()\n)). Please implement and register grad sampler for this layer. (See opacus.grad_sample.utils.register_grad_sampler)"), NotImplementedError("Model contains a trainable layer that Opacus doesn't currently support(blocks.3.0.conv2d1.conv2d:ParametrizedConv2d(\n 192, 192, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)\n (parametrizations): ModuleDict(\n (weight): ParametrizationList(\n (0): _SpectralNorm()\n )\n )\n)). Please implement and register grad sampler for this layer. (See opacus.grad_sample.utils.register_grad_sampler)"), NotImplementedError("Model contains a trainable layer that Opacus doesn't currently support(blocks.3.0.conv2d1.conv2d.parametrizations.weight:ParametrizationList(\n (0): _SpectralNorm()\n)). Please implement and register grad sampler for this layer. (See opacus.grad_sample.utils.register_grad_sampler)"), NotImplementedError("Model contains a trainable layer that Opacus doesn't currently support(blocks.3.0.conv2d2.conv2d:ParametrizedConv2d(\n 192, 192, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)\n (parametrizations): ModuleDict(\n (weight): ParametrizationList(\n (0): _SpectralNorm()\n )\n )\n)). Please implement and register grad sampler for this layer. (See opacus.grad_sample.utils.register_grad_sampler)"), NotImplementedError("Model contains a trainable layer that Opacus doesn't currently support(blocks.3.0.conv2d2.conv2d.parametrizations.weight:ParametrizationList(\n (0): _SpectralNorm()\n)). Please implement and register grad sampler for this layer. (See opacus.grad_sample.utils.register_grad_sampler)"), NotImplementedError("Model contains a trainable layer that Opacus doesn't currently support(blocks.4.0.conv2d1.conv2d:ParametrizedConv2d(\n 192, 192, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)\n (parametrizations): ModuleDict(\n (weight): ParametrizationList(\n (0): _SpectralNorm()\n )\n )\n)). Please implement and register grad sampler for this layer. (See opacus.grad_sample.utils.register_grad_sampler)"), NotImplementedError("Model contains a trainable layer that Opacus doesn't currently support(blocks.4.0.conv2d1.conv2d.parametrizations.weight:ParametrizationList(\n (0): _SpectralNorm()\n)). Please implement and register grad sampler for this layer. (See opacus.grad_sample.utils.register_grad_sampler)"), NotImplementedError("Model contains a trainable layer that Opacus doesn't currently support(blocks.4.0.conv2d2.conv2d:ParametrizedConv2d(\n 192, 192, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)\n (parametrizations): ModuleDict(\n (weight): ParametrizationList(\n (0): _SpectralNorm()\n )\n )\n)). Please implement and register grad sampler for this layer. (See opacus.grad_sample.utils.register_grad_sampler)"), NotImplementedError("Model contains a trainable layer that Opacus doesn't currently support(blocks.4.0.conv2d2.conv2d.parametrizations.weight:ParametrizationList(\n (0): _SpectralNorm()\n)). Please implement and register grad sampler for this layer. (See opacus.grad_sample.utils.register_grad_sampler)"), NotImplementedError("Model contains a trainable layer that Opacus doesn't currently support(linear1.liner:ParametrizedLinear(\n in_features=192, out_features=1, bias=True\n (parametrizations): ModuleDict(\n (weight): ParametrizationList(\n (0): _SpectralNorm()\n )\n )\n)). Please implement and register grad sampler for this layer. (See opacus.grad_sample.utils.register_grad_sampler)"), NotImplementedError("Model contains a trainable layer that Opacus doesn't currently support(linear1.liner.parametrizations.weight:ParametrizationList(\n (0): _SpectralNorm()\n)). Please implement and register grad sampler for this layer. (See opacus.grad_sample.utils.register_grad_sampler)")]

I customized a SpectralNorm based on the error message and wrapped it with opacus.grad_sample.utils.register_grad_sampler, which produces an error :

RuntimeError: mutating directly with .data inside functorch transform is not allowed.

The custom SpectralNorm comes from this python file.

[HuanyuZhang]

I have checked the implementation of "nn.utils.spectral_norm" and this. I think the main problem comes from their use of "_forward_pre_hooks", either explicitly or implicitly, which Opacus does not support right now. As a consequence, "weight_original" (or "w_bar" in the latter) is recognized as a trainable parameter but "weight" instead is used in the forward pass. This discrepancy confuses Opacus, leadibg to the empty per_sample_grad for "weight_original".

[HuanyuZhang]

Unfortunately, I do not have a quick solution in my mind right now (cc @karthikprasad in case you have a different opinion). And we might not have bandwidth to support this function in the very near future. Right now, I would suggest you just turn spectral norm off, or replace it by other normalization technique (like layer_normalization).