Difference between the GPU memory usage for MONAI UNet vs a custom UNet for a single forward pass through the networks

[ahxmeds]

Hi all,

I am using two UNets and comparing their GPU memory usage when I just run output = model(input). The first UNet is a custom UNet (UNetCustom) which contains about 1.3 million trainable parameters and is defined below. This UNet architecture comes from the following paper: https://ieeexplore.ieee.org/document/8581448

class DoubleConvolution(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(DoubleConvolution, self).__init__()
        self.double_conv_block = nn.Sequential(
            nn.Conv3d(in_channels, out_channels, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm3d(out_channels),
            nn.LeakyReLU(inplace=True),
            nn.Conv3d(out_channels, out_channels, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm3d(out_channels),
            nn.LeakyReLU(inplace=True)
        )
    def forward(self, x):
        return self.double_conv_block(x)

class DownsamplingOperation(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(DownsamplingOperation, self).__init__()
        self.downsampling_op = nn.Sequential(
            nn.Conv3d(in_channels, out_channels, kernel_size=3, stride=2, padding=1),
            nn.BatchNorm3d(out_channels),
            nn.LeakyReLU(inplace=True),
        )
    def forward(self, x):
        return self.downsampling_op(x)

class UpsamplingOperation(nn.Module):
    def __init__(self, scale_factor, in_channels, out_channels):
        super(UpsamplingOperation, self).__init__()

        self.upsampling_op = nn.Sequential(
            nn.Upsample(scale_factor=scale_factor, mode='trilinear', align_corners=True),
            nn.Conv3d(in_channels, out_channels, kernel_size=1),
            nn.BatchNorm3d(out_channels),
            nn.LeakyReLU(inplace=True)
        )
    def forward(self, x):
        return self.upsampling_op(x)
    

class UNetCustom(nn.Module):
    def __init__(self, in_channels=1, out_channels=1):
        super(UNetCustom, self).__init__()

        self.in_channels = in_channels
        self.out_channels = out_channels

        self.encoder1 = DoubleConvolution(self.in_channels, 16)
        self.downsample1 = DownsamplingOperation(16, 16)
        self.encoder2 = DoubleConvolution(16, 32)
        self.downsample2 = DownsamplingOperation(32, 32)
        self.encoder3 = DoubleConvolution(32, 64)
        self.downsample3 = DownsamplingOperation(64, 64)
        self.bottleneck = DoubleConvolution(64, 128)
        self.upsample3 = UpsamplingOperation(2, 128, 64)
        self.decoder3 = DoubleConvolution(64, 64)
        self.upsample2 = UpsamplingOperation(2, 64, 32)
        self.decoder2 = DoubleConvolution(32, 32)
        self.upsample1 = UpsamplingOperation(2, 32, 16)
        self.decoder1 = DoubleConvolution(16, self.out_channels)

    def forward(self, x):
        x1 = self.encoder1(x)
        x2 = self.downsample1(x1)
        x3 = self.encoder2(x2)
        x4 = self.downsample2(x3)
        x5 = self.encoder3(x4)
        x6 = self.downsample3(x5)
        x7 = self.bottleneck(x6)
        x8 = self.upsample3(x7) + x5
        x9 = self.decoder3(x8)
        x10 = self.upsample2(x9) + x3
        x11 = self.decoder2(x10)
        x12 = self.upsample1(x11) + x1 
        out = self.decoder1(x12)
        return out

unet_custom = UNetCustom(1,1).to(device)

The second UNet is the MONAI UNet from monai.networks.nets import UNet which contains about 19.3 million trainable parameters and is defined as follows:

from monai.networks.nets import UNet
from monai.networks.layers import Norm

unet_monai = UNet(
            spatial_dims=3,
            in_channels=1,
            out_channels=1,
            channels=(16, 32, 64, 128, 256, 512),
            strides=(2, 2, 2, 2, 2),
            num_res_units=2,
            norm=Norm.BATCH
        ).to(device)

I am using an NVIDIA V100 GPU with 16 GiB RAM. Even though the number of parameters are more for unet_monai as compared to unet_custom, when doing a forward pass using an input tensor of size [1, 1, 256, 256, 256], unet_custom takes up close to 15.5 GiB RAM, while unet_monai takes up only 9-10 GiB RAM. I am currently just testing this custom UNet and would like it to be able to use input of max input size [1, 1, 256, 256, 256].

What's the difference between the MONAI implementation of UNet vs. my custom implementation? Does anyone have any insights on where could this difference in memory usage be coming from?

Thanks in advance!

[KumoLiu]

Hi @ahxmeds, thanks for your interest here.
First I would suggest you use torchsummary to see the summary of the model.
MONAI implementation

Your implementation

You could find that your "Forward/backward pass size" is much larger than MONAI's.
What is meant by "Forward/backward pass size", you could refer to sksq96/pytorch-summary#51 (comment).

Finally, the main reason is that MONAI downsamples in the first conv.

Hope it helps, thanks!