Request: Extension for GroupNorm
- Dominant language
- Python
- Stars
- 617
- Forks
- 57
- PR merge metrics
- No merged PRs in 30d
Description
I am re-implementing the enhancement of DP-SGD through the [random sparsification](https://github.com/JunyiZhu-AI/RandomSparsification) of gradients on my UNet Model.
Here is a Debug info on extending the Model `extend(model)`.
```
UNet(
(encoder1): Sequential(
(enc1conv1): Conv2d(1, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(enc1norm1): GroupNorm(32, 32, eps=1e-05, affine=True)
(enc1relu1): ReLU()
(enc1conv2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(enc1norm2): GroupNorm(32, 32, eps=1e-05, affine=True)
(enc1relu2): ReLU()
)
(pool1): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
....
(bottleneck): Sequential(
(bottleneckconv1): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bottlenecknorm1): GroupNorm(32, 512, eps=1e-05, affine=True)
(bottleneckrelu1): ReLU()
(bottleneckconv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bottlenecknorm2): GroupNorm(32, 512, eps=1e-05, affine=True)
(bottleneckrelu2): ReLU()
)
(upconv4): ConvTranspose2d(512, 256, kernel_size=(2, 2), stride=(2, 2))
(decoder4): Sequential(
(dec4conv1): Conv2d(512, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(dec4norm1): GroupNorm(32, 256, eps=1e-05, affine=True)
(dec4relu1): ReLU()
(dec4conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(dec4norm2): GroupNorm(32, 256, eps=1e-05, affine=True)
(dec4relu2): ReLU()
)
...
)
```
BackPACK library does not support some of the modules in the model, specifically `**GroupNorm**`.
For this should I be creating custom extensions for the unsupported modules?
Logs when training the model.
```
env/lib/python3.11/site-packages/backpack/extensions/backprop_extension.py:106: UserWarning: Extension saving to grad_batch does not have an extension for Module although the module has parameters
warnings.warn(
env/lib/python3.11/site-packages/backpack/extensions/backprop_extension.py:106: UserWarning: Extension saving to grad_batch does not have an extension for Module although the module has parameters
```
Thanks for the help :)
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.