Pytorch model runtime error when testing U-Net
I've defined a U-Net model using Pytorch but it won't accept my input. I've checked the model layers and they seem to be applying the operations as I would expect them to but I still get an error.
I've just switched to Pytorch after mostly using Keras so I'm not really sure how to debug this issue, the error I get is:

RuntimeError: Given groups=1, weight of size [32, 64, 3, 3], expected
input[1, 128, 65, 65] to have 64 channels, but got 128 channels
instead

Here's the code I'm using:
class UNET(nn.Module):
    def __init__(self, in_channels=2, out_channels=2):
        super().__init__()

        self.conv1 = self.contract_block(in_channels, 32, 3, 1)
        self.conv2 = self.contract_block(32, 64, 3, 1)
        self.conv3 = self.contract_block(64, 128, 3, 1)

        self.upconv3 = self.expand_block(128, 64, 3, 1)
        self.upconv2 = self.expand_block(64, 32, 3, 1)
        self.upconv1 = self.expand_block(32, out_channels, 3, 1)

    def __call__(self, x):

        # downsampling part
        conv1 = self.conv1(x)
        conv2 = self.conv2(conv1)
        conv3 = self.conv3(conv2)

        upconv3 = self.upconv3(conv3)

        upconv2 = self.upconv2(torch.cat([upconv3, conv2], 1))
        upconv1 = self.upconv1(torch.cat([upconv2, conv1], 1))

        return upconv1

    def contract_block(self, in_channels, out_channels, kernel_size, padding):

        contract = nn.Sequential(
            torch.nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, stride=1, padding=padding),
            torch.nn.BatchNorm2d(out_channels),
            torch.nn.Tanh(),
            torch.nn.Conv2d(out_channels, out_channels, kernel_size=kernel_size, stride=1, padding=padding),
            torch.nn.BatchNorm2d(out_channels),
            torch.nn.Tanh(),
            torch.nn.MaxPool2d(kernel_size=2, stride=2, padding=1)
            )

        return contract

    def expand_block(self, in_channels, out_channels, kernel_size, padding):

        expand = nn.Sequential(
            torch.nn.Conv2d(in_channels, out_channels, kernel_size, stride=1, padding=padding),
            torch.nn.BatchNorm2d(out_channels),
            torch.nn.Tanh(),
            torch.nn.Conv2d(out_channels, out_channels, kernel_size, stride=1, padding=padding),
            torch.nn.BatchNorm2d(out_channels),
            torch.nn.Tanh(),
            torch.nn.ConvTranspose2d(out_channels, out_channels, kernel_size=2, stride=2, padding=1, output_padding=1)
            )
        return expand

model = UNET()
x = torch.randn(1, 2, 256, 256)
print(model(x))


Your problem is in the model layer definition.
You defined self.upconv2 = self.expand_block(64, 32, 3, 1) but what you do is concatenating 2 tensors each with 64 channels so in total you get 128.
You should fix the channels of the up-sampling part of the U-Net to match the number of channels after the concatenation.
Doing the mentioned fix will give you:
class UNET(nn.Module):
    def __init__(self, in_channels=2, out_channels=2):
        super().__init__()

        self.conv1 = self.contract_block(in_channels, 32, 3, 1)
        self.conv2 = self.contract_block(32, 64, 3, 1)
        self.conv3 = self.contract_block(64, 128, 3, 1)

        self.upconv3 = self.expand_block(128, 64, 3, 1)
        self.upconv2 = self.expand_block(64 + 64, 32, 3, 1)
        self.upconv1 = self.expand_block(32 + 32, out_channels, 3, 1)

    def __call__(self, x):

        # downsampling part
        conv1 = self.conv1(x)
        conv2 = self.conv2(conv1)
        conv3 = self.conv3(conv2)

        upconv3 = self.upconv3(conv3)
        upconv2 = self.upconv2(torch.cat([upconv3, conv2], 1))
        upconv1 = self.upconv1(torch.cat([upconv2, conv1], 1))

        return upconv1

    def contract_block(self, in_channels, out_channels, kernel_size, padding):

        contract = nn.Sequential(
            torch.nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, stride=1, padding=padding),
            torch.nn.BatchNorm2d(out_channels),
            torch.nn.Tanh(),
            torch.nn.Conv2d(out_channels, out_channels, kernel_size=kernel_size, stride=1, padding=padding),
            torch.nn.BatchNorm2d(out_channels),
            torch.nn.Tanh(),
            torch.nn.MaxPool2d(kernel_size=2, stride=2, padding=1)
            )

        return contract

    def expand_block(self, in_channels, out_channels, kernel_size, padding):

        expand = nn.Sequential(
            torch.nn.Conv2d(in_channels, out_channels, kernel_size, stride=1, padding=padding),
            torch.nn.BatchNorm2d(out_channels),
            torch.nn.Tanh(),
            torch.nn.Conv2d(out_channels, out_channels, kernel_size, stride=1, padding=padding),
            torch.nn.BatchNorm2d(out_channels),
            torch.nn.Tanh(),
            torch.nn.ConvTranspose2d(out_channels, out_channels, kernel_size=2, stride=2, padding=1, output_padding=1)
            )
        return expand

model = UNET()
x = torch.randn(1, 2, 256, 256)
print(model(x))


Based on the comment, if you have other spatial dimensions that might not fit the parameters of the convolutions operations you can do one of 2 options:

Start play with the parameter based on the formula at the bottom of the Conv2d so that you will match the input dimension.

You could force pad the target to desired dimension using the following 2 functions:
import torch
import torch.nn as nn
import torch.nn.functional as F


def pad_tensor(source, target):
    &quot;&quot;&quot;
    Pad source tensor to match target tensor size

    :param source: tensor that need to get padding
    :param target: tensor of the desired shape
    :return: source tensor with shape equal to target
    &quot;&quot;&quot;

    diff_y = target.size()[2] - source.size()[2]
    diff_x = target.size()[3] - source.size()[3]

    source = F.pad(source, [diff_x // 2, diff_x - diff_x // 2, diff_y // 2, diff_y - diff_y // 2])
    return source


def concatenate_tensors(x1, x2):
    &quot;&quot;&quot;
    Concatenate both tensors

    :param x1: first tensor to be concatenated
    :param x2: second tensor to be concatenated
    :return: concatenation of both tensors
    &quot;&quot;&quot;

    x1 = pad_tensor(x1, x2)
    return torch.cat([x1, x2], dim=1)



Now call concatenate_tensors instead of torch.cat this will fix the dimensions to match the size you need.
