Train two UNet in a pipeline

[J_Johnson]

Just to add to what @KFrank said, here is a code implementation combining two model’s parameters within an optimizer:

Merging two models

You could just add the parameters lists: optimizer = optim.SGD(list(modelA.parameters()) + list(modelB.parameters()), lr=1e-3) How are you transferring the parameters from layer A2 to B1? If so, the weight matrix will have a size mismatch ([30, 8] vs. [40, 8]).

Alternatively, you can wrap both models in an nn.Module class and just pass that into an optimizer. That might look something like the following:

class SuperModel(nn.Module):
    def __init__(self, n1, n2):
        super(SuperModel, self).__init__()

        model1 = UNet(in_channels = 3, out_channels = 1, hidden_size = n1)
        model2 = UNet(in_channels = 4, out_channels = 1, hidden_size = n2)

    def forward(self, x):
        mid_output = model1(x)
        
        mid_input = torch.cat([x, mid_output], dim = 1)
        final_output = model2(mid_input)
        return final_output


model = SuperModel(320, 320)

optimizer = torch.optim.SGD(model.parameters(), lr = 0.01)