Multitask unet for dual segmentation

Claudia_22: <p>is it possible to train one unet with a shared encoder and dual decoders to perform two different tasks?</p>

KFrank: <p>Hi Claudia!</p>
<aside class="quote no-group" data-username="Claudia_22" data-post="1" data-topic="220619" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/c/c67d28/48.png" class="avatar"> Claudia_22:</div>
<blockquote>
<p>is it possible to train one unet with a shared encoder and dual decoders to perform two different tasks?</p>
</blockquote>
</aside>
<p>I’ve never tried it and I don’t know off hand of any examples or research, but I think that<br>
this would be possible.</p>
<p>Some commentary:</p>
<p>First, for simplicity, let’s assume that the input for a baseline, single-task U-Net would be a<br>
single-channel gray-scale image and that the output would be a single-channel probabilistic<br>
binary (e.g., <em>logits</em>) prediction “image.”</p>
<p>The simplest approach would be to take your “regular” single-task U-Net and change its<br>
final layer to have two output channels (one for task A and the second for task B), instead<br>
of only one.</p>
<p>In this case almost all of the U-Net is shared between, and jointly trained on, the two tasks,<br>
with only each of the two output channels specific to a given task.</p>
<p>In your proposal, the first half of the U-Net – the “encoder” – is shared and trained jointly,<br>
while the second half – the dual “decoders” – are task specific and trained separately.</p>
<p>The next logical step is to realize that you could switch from joint to separate anywhere<br>
along the U-Net.  For example, it would be logically sound to have the first two layers of<br>
the U-Net’s encoder be shared, but have most of the encoder (as well as the decoder)<br>
be dual.</p>
<p>If you do decide to experiment with a joint-encoder, dual-decoder architecture, it would be<br>
good practice to start with a “regular” two-channel-output U-Net as a baseline for comparison.</p>
<p>(You should also clarify in your own mind the details of your use case.  For example, are you<br>
predicting individual pixels to be either task A or task B or neither, so a three-class per-pixel<br>
multi-class classification problem, or would you be making two binary predictions for each pixel<br>
separately for each of task A and task B, so a two-class <em>multi-label</em> classification problem?)</p>
<p>Best.</p>
<p>K. Frank</p>

Claudia_22: <p>Thank you Frank. So the thing is , one of the decoders will give a binary  class output , and the other will give a multiclass output. however , the encoder extracts features from the same data space. Is that also possible?</p>

Claudia_22: <p>And also will the combination of the loss functions that were dedicated to each of the decoders be used to update the weights in the encoder ?</p>

KFrank: <p>Hi Claudia!</p>
<aside class="quote no-group" data-username="Claudia_22" data-post="3" data-topic="220619" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/c/c67d28/48.png" class="avatar"> Claudia_22:</div>
<blockquote>
<p>one of the decoders will give a binary  class output , and the other will give a multiclass output. however , the encoder extracts features from the same data space. Is that also possible?</p>
</blockquote>
</aside>
<p>Yes.  The first encoder would have a single-channel output – the logit for the binary-classification<br>
prediction.  The second encoder would have a multi-channel output.</p>
<p>Suppose that the second encoder is classifying pixels for semantic segmentation, let’s say<br>
for “hat,” “shirt,”, “pants,” “shoes,” and “background.”  Then the second decoder would have<br>
a five-channel output – the unnormalized log-probabilities for the five segmentation classes.</p>
<aside class="quote no-group" data-username="Claudia_22" data-post="4" data-topic="220619" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/c/c67d28/48.png" class="avatar"> Claudia_22:</div>
<blockquote>
<p>And also will the combination of the loss functions that were dedicated to each of the decoders be used to update the weights in the encoder ?</p>
</blockquote>
</aside>
<p>In the baseline approach of using a “regular” U-Net that I mentioned in an earlier post, the last<br>
layer of the (single-decoder) U-Net would have six channels.  You would take the first channel<br>
and pass it to <code>BCEWithLogitsLoss</code> (for the binary segmentation) and pass the remaining five<br>
channels, grouped together, to <code>CrossEntropyLoss</code> (for the multi-class segmentation).</p>
<p>With two separate decoders, you would do essentially the same thing – use the output of the<br>
single-channel decoder for binary segmentation and the output of the five-channel decoder<br>
for multi-class segmentation.</p>
<p>You would sum the two losses together, probably with weights, into a total loss, and then<br>
backpropagate the total loss.  The gradients computed for, say, the single-channel decoder<br>
will depend only on the <code>BCEWithLogitsLoss</code>, while the gradients computed for the encoder<br>
will depend on both the <code>BCEWithLogitsLoss</code> and the <code>CrossEntropyLoss</code>, so the encoder<br>
will be trained to produce an encoding that is useful for both segmentation tasks.</p>
<p>Again, don’t be fixated on having two separate decoders.  There may well be advantages to<br>
having essentially all of the U-Net be shared between – and trained jointly on – both of the<br>
segmentation tasks.  You should definitely try the single-joint-decoder approach and only<br>
switch over to the dual-decoder architecture  if it is empirically seen to work better.</p>
<p>(You could argue that such a joint decoder “is doing more work” than one of the single-task<br>
decoders.  So you could argue that joint-decoder architecture should have more internal<br>
channels than either of the single-task decoders does individually.  To make a fair comparison<br>
between the two architectures, you would have to experiment with that.)</p>
<p>Best.</p>
<p>K. Frank</p>

Claudia_22: <p>Thank you very much Frank. I think i wasn’t clear enough on the kind of task i am doing.<br>
Let’s say i have a dataset which has MRI data of the prostate gland. I have one groundtruth label mask for tumours (two labels) and another groundtruth label mask for the different divisions of the prostate gland (multiple labels). I want to perform these two segmentation tasks simultaneously using a UNET with a shared encoder with dual decoders whereby each decoder performs a specific segmentation task due to the two different groundtruth masks.<br>
Based on everything that has been discussed is it possible to d something like that?</p>