Train two UNet in a pipeline

vvuonghn: <p>Hello everyone,</p>
<p>I am implementing the segmentation as the pipeline below. How can I concatenate 2 Unet as the same pipeline, and training together (without training 2 models separately)? The training/inference is the same pipeline.</p>
<p>For loss function: I have 3 loss function for Unet1 (loss_1), Unet2 (loss_2), loss = loss_1 + loss_2</p>
<p>Thank you.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/5/6/56105563dad47d332906015d5718cbaaa96de77c.png" data-download-href="https://discuss.pytorch.org/uploads/default/56105563dad47d332906015d5718cbaaa96de77c" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/5/6/56105563dad47d332906015d5718cbaaa96de77c_2_690x259.png" alt="image" data-base62-sha1="chm3wS4ifoCmeKEY15BSzzSs56k" width="690" height="259" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/5/6/56105563dad47d332906015d5718cbaaa96de77c_2_690x259.png, https://discuss.pytorch.org/uploads/default/optimized/3X/5/6/56105563dad47d332906015d5718cbaaa96de77c_2_1035x388.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/5/6/56105563dad47d332906015d5718cbaaa96de77c_2_1380x518.png 2x" data-dominant-color="F4F4F4"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1459×548 38.1 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>

KFrank: <p>Hi Ngoc!</p>
<aside class="quote no-group" data-username="vvuonghn" data-post="1" data-topic="188310" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/v/13edae/48.png" class="avatar"> vvuonghn:</div>
<blockquote>
<p>How can I concatenate 2 Unet as the same pipeline, and training together (without training 2 models separately)?</p>
</blockquote>
</aside>
<p>I’m not sure that I understand your use case, but perhaps you are asking for<br>
something like the following:</p>
<p>Let me assume that the input to your first U-Net is a three-channel (RGB) image,<br>
and that its output is a single “probability” per pixel, so one channel.  You want<br>
to take the output of the first U-Net an use it, together with another RGB image<br>
(or perhaps the same RGB image), as input to the second U-Net.  Last, let me<br>
assume that you want the “predictions” output by the second U-Net to be a single<br>
value per pixel.</p>
<p>Then your first U-Net should have three input channels and one output channel.<br>
Your second U-Net should have, however, <em>four</em> input channels (and one output<br>
channel).  If your second U-Net is based on a three-input-channel U-Net, you will<br>
have to modify its first convolutional layer to use four input channels, rather than<br>
three.</p>
<p>Let <code>output1</code> have shape <code>[nBatch, 1, height, width]</code> and be the output of<br>
the first U-Net.  Let the second <code>RGB2</code> have shape <code>[nBatch, 3, height, width]</code>.<br>
Then <code>torch.cat ((output1, RGB2), dim = 1)</code> will have four channels (thus<br>
shape <code>[nBatch, 4, height, width]</code>) and be suitable as input to the second<br>
four-input-channel U-Net.</p>
<p>You can certainly backpropagate using a combined loss that is the sum of two<br>
terms, one that depends only on the output of the first U-Net and a second that<br>
depends on the output of the second U-Net.  Doing so will populate the gradients<br>
of the parameters of the two U-Nets in the way I believe you want.  You can then<br>
run an optimizer that contains both U-Nets’ parameters in its parameter list.</p>
<p>Best.</p>
<p>K. Frank</p>

J_Johnson: <p>Just to add to what <a class="mention" href="/u/kfrank">@KFrank</a> said, here is a code implementation combining two model’s parameters within an optimizer:</p>
<aside class="quote" data-post="2" data-topic="45637">
  <div class="title">
    <div class="quote-controls"></div>
    <img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/48/1823_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/merging-two-models/45637/2">Merging two models</a> <a class="badge-category__wrapper " href="/c/vision/5"><span data-category-id="5" data-drop-close="true" class="badge-category " title="Topics related to either pytorch/vision or vision research related topics"><span class="badge-category__name">vision</span></span></a>
  </div>
  <blockquote>
    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]).
  </blockquote>
</aside>

<p>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:</p>
<pre><code class="lang-auto">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)
</code></pre>

vvuonghn: <p>Thank you everyone for your support.<br>
Based on your suggestion, I can implement the pipeline and train the model</p>