Upsampling ResBlock GPU memory spike
Labels: module: cuda, module: memory usage, triaged

### 🐛 Describe the bug

During inference, my PyTorch model experiences a significant GPU memory spike of up to 10GB with the first batch only, while the following batches’ memory consumption caps at 3GB. I have identified that the ResBlock upsampling layers in my model’s architecture are causing the spike. I have six upsampling ResBlock layers, and the spike only occurs when the bottom three ResBlock layers are included. The problematic layers are Resblock [2-51, 2-52, 2-53] from the layers listed below. When I increase the batch size, the spike lasts longer, and instead of spiking with the first batch only, it can stay until the end of all batches.

**Layers’ Summary:**

```
└─ResBlock: 2-48                              [2, 512, 8, 8]            1
   │    └─Sequential: 3-71                       [2, 512, 8, 8]            4,718,592
   │    └─Sequential: 3-72                       [2, 512, 8, 8]            262,656
└─ResBlock: 2-49                              [2, 512, 16, 16]          1
   │    └─Sequential: 3-73                       [2, 512, 16, 16]          4,718,592
   │    └─Sequential: 3-74                       [2, 512, 16, 16]          262,656
└─ResBlock: 2-50                              [2, 512, 32, 32]          1
   │    └─Sequential: 3-75                       [2, 512, 32, 32]          4,718,592
   │    └─Sequential: 3-76                       [2, 512, 32, 32]          262,656
└─ResBlock: 2-51                              [2, 256, 64, 64]          1
   │    └─Sequential: 3-77                       [2, 256, 64, 64]          1,769,472
   │    └─Sequential: 3-78                       [2, 256, 64, 64]          131,328
└─ResBlock: 2-52                              [2, 128, 128, 128]        1
   │    └─Sequential: 3-79                       [2, 128, 128, 128]        442,368
   │    └─Sequential: 3-80                       [2, 128, 128, 128]        32,896
└─ResBlock: 2-53                              [2, 64, 256, 256]         1
   │    └─Sequential: 3-81                       [2, 64, 256, 256]         110,592
   │    └─Sequential: 3-82                       [2, 64, 256, 256]         8,256
```

**ResBlock Code:**

```python
class ResBlock(nn.Module):
    def __init__(self, in_channels, out_channels, padding, upsample, downsample,
                 norm_layer, activation=nn.ReLU, gated=False):
        super(ResBlock, self).__init__()
        normalize = True
        bias = False
        self.noise_strength = torch.nn.Parameter(torch.zeros([]))
        self.upsample = True
        norm0 = AdaptiveNorm2d(in_channels, norm_layer)
        norm1 = AdaptiveNorm2d(out_channels, norm_layer)

        layers = []
        layers.append(norm0)
        layers.append(activation(inplace=True))
        layers.append(nn.Upsample(scale_factor=2, mode='bilinear'))

        layers.extend([
            nn.Sequential() if padding is nn.ZeroPad2d else padding(1),
            spectral_norm(
                nn.Conv2d(in_channels, out_channels, 3, 1, 1 if padding is nn.ZeroPad2d else 0, bias=bias),
                eps=1e-4)])
        layers.append(norm1)
        layers.extend([
            activation(inplace=True),
            nn.Sequential() if padding is nn.ZeroPad2d else padding(1),
            spectral_norm(
                nn.Conv2d(out_channels, out_channels, 3, 1, 1 if padding is nn.ZeroPad2d else 0, bias=bias),
                eps=1e-4)])
        self.block = nn.Sequential(*layers)

        layers = []
        layers.append(nn.Upsample(scale_factor=2, mode='bilinear'))
        layers.append(spectral_norm(
        nn.Conv2d(in_channels, out_channels, 1), eps=1e-4))
        self.skip = nn.Sequential(*layers)

    def forward(self, input):
        out = self.block(input)
        self.noise_strength.requires_grad = False
        output = out + self.skip(input)
        return output
```
I would appreciate suggestions to overcome the memory spike issue without compromising the ResBlock upsampling layers.

### Versions

PyTorch version: 1.12.1+cu113
Is debug build: False
CUDA used to build PyTorch: 11.3
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.6 LTS (x86_64)
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
Clang version: Could not collect
CMake version: version 3.22.3
Libc version: glibc-2.27

Python version: 3.10.9 (main, Jan 28 2023, 19:09:35) [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-5.4.0-1068-aws-x86_64-with-glibc2.27
Is CUDA available: True
CUDA runtime version: 9.1.85
CUDA_MODULE_LOADING set to: 
GPU models and configuration: GPU 0: NVIDIA A10G
Nvidia driver version: 470.82.01
cuDNN version: Probably one of the following:
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.0.5
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn.so.8.1.1
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.1.1
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.1.1
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.1.1
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.1.1
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.1.1
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.1.1
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:        x86_64
CPU op-mode(s):      32-bit, 64-bit
Byte Order:          Little Endian
CPU(s):              16
On-line CPU(s) list: 0-15
Thread(s) per core:  2
Core(s) per socket:  8
Socket(s):           1
NUMA node(s):        1
Vendor ID:           AuthenticAMD
CPU family:          23
Model:               49
Model name:          AMD EPYC 7R32
Stepping:            0
CPU MHz:             2729.895
BogoMIPS:            5600.00
Hypervisor vendor:   KVM
Virtualization type: full
L1d cache:           32K
L1i cache:           32K
L2 cache:            512K
L3 cache:            16384K
NUMA node0 CPU(s):   0-15
Flags:               fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good nopl nonstop_tsc cpuid extd_apicid aperfmperf tsc_known_freq pni pclmulqdq ssse3 fma cx16 sse4_1 sse4_2 movbe popcnt aes xsave avx f16c rdrand hypervisor lahf_lm cmp_legacy cr8_legacy abm sse4a misalignsse 3dnowprefetch topoext ssbd ibrs ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 clzero xsaveerptr wbnoinvd arat npt nrip_save rdpid

Versions of relevant libraries:
[pip3] efficientnet-pytorch==0.6.3
[pip3] numpy==1.23.5
[pip3] segmentation-models-pytorch==0.2.1
[pip3] torch==1.12.1+cu113
[pip3] torchaudio==0.12.1+cu113
[pip3] torchinfo==1.7.2
[pip3] torchstat==0.0.7
[pip3] torchsummary==1.5.1
[pip3] torchvision==0.13.1+cu113
[conda] blas                      1.0                         mkl  
[conda] mkl                       2020.2                      256  
[conda] mkl-service               2.3.0            py37he8ac12f_0  
[conda] mkl_fft                   1.2.1            py37h54f3939_0  
[conda] mkl_random                1.1.1            py37h0573a6f_0  
[conda] numpy                     1.21.5                   pypi_0    pypi
[conda] numpydoc                  1.1.0              pyhd3eb1b0_1  

cc @ngimel