Cuda out of memory in 3D U-Net

ray_chiu: <p>Hello everyone.</p>
<p>I wondered if anyone else out there was using 3D U-Net in Pytorch and having trouble with Cuda out of memory issue? I’m trying to train a 3D U-Net model on Colab pro (with GPU memory 16GB) to predict 2 classes from 3D medical image with 512<em>512</em>N in size and keep facing cuda out of memory issue.<br>
I’ve try torch.cuda.empty_cache() but doesn’t work.</p>
<p>My U-Net architecture look like this:</p>
<p>class DoubleConv(nn.Module):<br>
“”"(Conv3D → BN → ReLU) * 2"""<br>
def <strong>init</strong>(self, in_channels, out_channels, num_groups=1):<br>
super().<strong>init</strong>()<br>
self.double_conv = nn.Sequential(<br>
nn.Conv3d(in_channels, out_channels, kernel_size=3, stride=1, padding=1),<br>
<span class="hashtag">#nn</span>.BatchNorm3d(out_channels),<br>
nn.GroupNorm(num_groups=num_groups, num_channels=out_channels),<br>
nn.ReLU(inplace=True),</p>
<pre><code>        nn.Conv3d(out_channels, out_channels, kernel_size=3, stride=1, padding=1),
        #nn.BatchNorm3d(out_channels),
        nn.GroupNorm(num_groups=num_groups, num_channels=out_channels),
        nn.ReLU(inplace=True)
      )

def forward(self,x):
    return self.double_conv(x)
</code></pre>
<p>class Down(nn.Module):</p>
<pre><code>def __init__(self, in_channels, out_channels):
    super().__init__()
    self.encoder = nn.Sequential(
        nn.MaxPool3d(2, 2),
        DoubleConv(in_channels, out_channels)
    )
def forward(self, x):
    return self.encoder(x)
</code></pre>
<p>class Up(nn.Module):</p>
<pre><code>def __init__(self, in_channels, out_channels, trilinear=True):
    super().__init__()
    
    if trilinear:
        self.up = nn.Upsample(scale_factor=2, mode='trilinear', align_corners=True)
    else:
        self.up = nn.ConvTranspose3d(in_channels // 2, in_channels // 2, kernel_size=2, stride=2)
        
    self.conv = DoubleConv(in_channels, out_channels)

def forward(self, x1, x2):
    x1 = self.up(x1)

    diffZ = x2.size()[2] - x1.size()[2]
    diffY = x2.size()[3] - x1.size()[3]
    diffX = x2.size()[4] - x1.size()[4]
    x1 = torch.nn.functional.pad(x1, [diffX // 2, diffX - diffX // 2, diffY // 2, diffY - diffY // 2, diffZ // 2, diffZ - diffZ // 2])

    x = torch.cat([x2, x1], dim=1)
    return self.conv(x)
</code></pre>
<p>class Out(nn.Module):<br>
def <strong>init</strong>(self, in_channels, out_channels):<br>
super().<strong>init</strong>()<br>
self.conv = nn.Conv3d(in_channels, out_channels, kernel_size = 1)</p>
<pre><code>def forward(self, x):
    return self.conv(x)
</code></pre>
<p>class UNet3d(nn.Module):<br>
def <strong>init</strong>(self, in_channels, n_classes, n_channels):<br>
super().<strong>init</strong>()<br>
self.in_channels = in_channels<br>
self.n_classes = n_classes<br>
self.n_channels = n_channels</p>
<pre><code>    self.conv = DoubleConv(in_channels, n_channels)
    self.enc1 = Down(n_channels, 2 * n_channels)
    self.enc2 = Down(2 * n_channels, 4 * n_channels)
    self.enc3 = Down(4 * n_channels, 8 * n_channels)
    self.enc4 = Down(8 * n_channels, 8 * n_channels)

    self.dec1 = Up(16 * n_channels, 4 * n_channels)
    self.dec2 = Up(8 * n_channels, 2 * n_channels)
    self.dec3 = Up(4 * n_channels, n_channels)
    self.dec4 = Up(2 * n_channels, n_channels)
    self.out = Out(n_channels, n_classes)

def forward(self, x):
    x1 = self.conv(x)
    x2 = self.enc1(x1)
    x3 = self.enc2(x2)
    x4 = self.enc3(x3)
    x5 = self.enc4(x4)

    mask = self.dec1(x5, x4)
    mask = self.dec2(mask, x3)
    mask = self.dec3(mask, x2)
    mask = self.dec4(mask, x1)
    mask = self.out(mask)
    return mask
</code></pre>
<p><span class="hashtag">#I</span> create a random tensor and the error still exist<br>
tensor = torch.randn(1,1,512,512,64)<br>
NET = UNet3d(in_channels=1, n_classes=2, n_channels=1).to(device)<br>
input = tensor.cuda()<br>
NET(input)</p>
<p>RuntimeError: CUDA out of memory. Tried to allocate 64.00 MiB (GPU 0; 14.76 GiB total capacity; 13.42 GiB already allocated; 47.75 MiB free; 13.48 GiB reserved in total by PyTorch) If reserved memory is &gt;&gt; allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF</p>

ptrblck: <p>You could try to decrease the spatial size of your model, decrease the number of parameters of your model (e.g. less layers or less filters/weights etc.), use <code>torch.utils.checkpoint</code> to trade compute for memory, or maybe also CPU offloading of activations.</p>

ray_chiu: <p>Thank you. I decreased the channel size of the model and it worked!</p>