Giving two inputs to as single model and concatenating them

Question

I have a model that takes one image as input and it works fine. now I want to give one more transformed image with the same dimensions as the first one as input to the model. The model should learn from both images. The below code shows an error: "init() takes 2 positional arguments but 3 were given". I want to know why I have to give two inputs in the init function as well.

 class MyNet(nn.Module):
    def __init__(self, input_dim):
        super(MyNet, self).__init__()
        self.conv1 = nn.Conv2d(input_dim, nChannel, kernel_size=3, stride=1, padding=1)
        self.bn1 = nn.BatchNorm2d(nChannel)
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = []
        self.bn2 = []
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        for i in range(nConv - 1):
            self.conv2.append(nn.Conv2d(nChannel, nChannel, kernel_size=3, stride=1, padding=1))
            self.bn2.append(nn.BatchNorm2d(nChannel))

        self.conv3 = nn.Conv2d(nChannel, nChannel, kernel_size=1, stride=1, padding=0)
        self.bn3 = nn.BatchNorm2d(nChannel)
        self.UB1 = nn.UpsamplingBilinear2d(scale_factor=2)
        self.deconv = nn.ConvTranspose2d(nChannel, nChannel, kernel_size=3, stride=1, padding=1)

    def forward(self, x1, x2):
        x1 = self.conv1(x1)
        x1 = F.relu(x1)
        x1 = self.bn1(x1)
        x1 = self.pool1(x1)
        for i in range(nConv - 1):
            x1 = self.conv2[i](x1)
            x1 = F.relu(x1)
            x1 = self.bn2[i](x1)
            if i == 0:
                x1 = self.pool2(x1)
        x1 = self.conv3(x1)
        # x = F.relu(x)
        x1 = self.bn3(x1)

        x1 = self.UB1(x1)
        x1 = self.deconv(x1)
        x1 = F.relu(x1)
        x1 = self.bn3(x1)

        x1 = self.UB1(x1)
        x1 = self.deconv(x1)
        x1 = F.relu(x1)
        x1 = self.bn3(x1)

    


        x2 = self.conv1(x2)
        x2 = F.relu(x2)
        x2 = self.bn1(x2)
        x2 = self.pool1(x2)
        for i in range(nConv - 1):
            x2 = self.conv2[i](x2)
            x2 = F.relu(x2)
            x2 = self.bn2[i](x2)
            if i == 0:
                x2 = self.pool2(x2)
        x2 = self.conv3(x2)
        # x = F.relu(x)
        x2 = self.bn3(x2)

        x2 = self.UB1(x2)
        x2 = self.deconv(x2)
        x2 = F.relu(x2)
        x2 = self.bn3(x2)

        x2 = self.UB1(x2)
        x2 = self.deconv(x2)
        x2 = F.relu(x2)
        x2 = self.bn3(x2)
    

        x = torch.cat(x1,x2)
        return x

Answer 1

You've mentioned that 2 input images have the same size, so no need to pass 2 values on model instantiating (model = MyNet(input_dim)).

But forward then requires 2 inputs in your case it looks ok

# init passing only single value `input_dim`
model = MyNet(input_dim)
...
# pseudo code for training stage
for batch in data_loader:
    img1, img2 = batch
    output = model(img1, img2)