Strange Error with MaxPool2D - C++ Issue

vision
1

I’m encountering a pretty strage error with PyTorch, “max_pool2d_with_indices_out_cuda_frame”. I’m using the PyTorch UNET implementation found at https://github.com/milesial/Pytorch-UNet. This error seems to be occurring after the 3rd Down Conv i.e self.down3. The input shape to that layer is [20, 256, 32, 32].

I don’t really seem to understand why the error occurs or what I can do to fix it. I’ve attached the relevant parts of the code below.

The forward code:

def forward(self, x):
        x1 = self.inc(x)
        x2 = self.down1(x1)
        x3 = self.down2(x2)
        x4 = self.down3(x3)
        x5 = self.down4(x4)
        x = self.up1(x5, x4)
        x = self.up2(x, x3)
        x = self.up3(x, x2)
        x = self.up4(x, x1)
        logits = self.outc(x)
        return self.sigmoid(logits)

The self.down3 module (its of the class type Down):

class DoubleConv(nn.Module):
    """(convolution => [BN] => ReLU) * 2"""

    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.double_conv = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        return self.double_conv(x)


class Down(nn.Module):
    """Downscaling with maxpool then double conv"""

    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.maxpool_conv = nn.Sequential(
            nn.MaxPool2d(2),
            DoubleConv(in_channels, out_channels)
        )

    def forward(self, x):
        return self.maxpool_conv(x)
2

This seems to be a bug with the current PyTorch version i.e 1.4.0. Downgrading to 1.2.0 fixes the issue for me

3

Could you post the error message please?

4

This is the error that I’m getting. Here’s a full version of the error

Traceback (most recent call last):
  File "train_new.py", line 197, in <module>
    train_net(net=net,
  File "train_new.py", line 90, in train_net
    pred_rgb = net(bayer_input)
  File "/home/ybelhe/anaconda3/envs/pytorch_p3/lib/python3.8/site-packages/torch/nn/modules/module.py", line 532, in __call__
    result = self.forward(*input, **kwargs)
  File "/home/ybelhe/code/training/unet/unet_model.py", line 51, in forward
    x4 = self.down3(x3)
  File "/home/ybelhe/anaconda3/envs/pytorch_p3/lib/python3.8/site-packages/torch/nn/modules/module.py", line 532, in __call__
    result = self.forward(*input, **kwargs)
  File "/home/ybelhe/code/training/unet/unet_parts.py", line 39, in forward
    return self.maxpool_conv(x)
  File "/home/ybelhe/anaconda3/envs/pytorch_p3/lib/python3.8/site-packages/torch/nn/modules/module.py", line 532, in __call__
    result = self.forward(*input, **kwargs)
  File "/home/ybelhe/anaconda3/envs/pytorch_p3/lib/python3.8/site-packages/torch/nn/modules/container.py", line 100, in forward
    input = module(input)
  File "/home/ybelhe/anaconda3/envs/pytorch_p3/lib/python3.8/site-packages/torch/nn/modules/module.py", line 532, in __call__
    result = self.forward(*input, **kwargs)
  File "/home/ybelhe/anaconda3/envs/pytorch_p3/lib/python3.8/site-packages/torch/nn/modules/pooling.py", line 139, in forward
    return F.max_pool2d(input, self.kernel_size, self.stride,
  File "/home/ybelhe/anaconda3/envs/pytorch_p3/lib/python3.8/site-packages/torch/_jit_internal.py", line 181, in fn
    return if_false(*args, **kwargs)
  File "/home/ybelhe/anaconda3/envs/pytorch_p3/lib/python3.8/site-packages/torch/nn/functional.py", line 487, in _max_pool2d
    return torch.max_pool2d(
RuntimeError: max_pool2d_with_indices_out_cuda_frame failed with error code 0

Please do let me know if there’s any other information that I can provide you with