How to Jointly Train Models that Split? [Advanced Computer Vision Paper]

username: dylan
cooked: <p>I’m interested in implementing the Aux loss model described in <a href="https://arxiv.org/pdf/2002.02582.pdf" rel="noopener nofollow ugc">this paper</a>. I have attached an image of the model architecture below.</p>
<p><strong>A small description of the model and how it works:</strong><br>
This model is used to classify X-ray images. PA and L are two different views of an X-ray for one patient. Specifically, PA is an X-ray viewed from the back and L the same X-ray but from a side view.</p>
<p>The paper leverages DenseNet121 pretrained models (4 dense blocks) for this architecture. In the image below, a circle represents a dense block.</p>
<p><strong>Here’s the logic/control flow of the model:</strong></p>
<ul>
<li>If the L image is present but the PA image is missing for a patient, the data should only flow through the L branch (i.e. the DenseNet121 model for L).</li>
<li>If the PA image is present but the L image is missing for a patient, the data should only flow through the PA branch (i.e. the DenseNet121 model for PA).</li>
<li>If both images are present, the data should flow through both branches up to the 3rd dense block. After that, the outputs of each of the 3rd dense blocks should be concatenated and passed through another model (assume it’s linear).<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/4/8/485f68f2dcf4f4f678f64da5feeb0da16ec7b0b6.png" data-download-href="https://discuss.pytorch.org/uploads/default/485f68f2dcf4f4f678f64da5feeb0da16ec7b0b6" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/4/8/485f68f2dcf4f4f678f64da5feeb0da16ec7b0b6_2_332x500.png" alt="image" data-base62-sha1="akeN9mShSVxn1CSKxXfQrzPaCN0" width="332" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/4/8/485f68f2dcf4f4f678f64da5feeb0da16ec7b0b6_2_332x500.png, https://discuss.pytorch.org/uploads/default/original/3X/4/8/485f68f2dcf4f4f678f64da5feeb0da16ec7b0b6.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/4/8/485f68f2dcf4f4f678f64da5feeb0da16ec7b0b6.png 2x" data-dominant-color="F3F4F5"><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">340×512 21.3 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></li>
</ul>
<p>I understand how I could implement and train this model jointly for a batch size of 1. However, the paper claims that they can train this model jointly for a batch size of 8. How could I implement this efficiently in PyTorch for a batch size of 8?</p>

username: JuanFMontesinos
cooked: <p>What do you mean by efficiently? It’s a matter of setting if statements in the forward pass. You “may” optimize it by using two optimizers to run backward only when needed.</p>

username: dylan
cooked: <p>Here’s some pseudo code for the forward pass I have  in my model so far:</p>
<pre><code class="lang-python">def forward(self, x):
   if x has PA view but not L view:
        go through PA classifier
   if x has L view but not PA view:
        go through L classifier
   else:
       output1 = run through the first 3 blocks of of PA classifier
       output2 = run through second 3 blocks of L classifier
       concatenate output1 and output2 and feed them through Linear Layer

   return output
</code></pre>
<p>Is the forward pass designed to assume a batch size of 1?</p>
<p>If not, wouldn’t I need to do this:</p>
<pre><code class="lang-python">def forward(self, x):
  for example in x:
      if x has PA view but not L view:
         .....

       return output
</code></pre>

username: JuanFMontesinos
cooked: <p>By-default Pytorch nn.Modules requires the first dimension to be the batch.<br>
In case of an RGB image a single sample is of the size [R,G,B]xW,H. The input by default requires an additional dimension so that it becomes Nx3xWxH.<br>
In general any model is designed this way. Increasing your batch size is nothing but making N=8.<br>
This concatanation step is automatically done by the dataloader. If you are using pytorch dataloader you just need to change the batch_size argument.</p>
<p>WRT optimizers, optimizers iterate over model parameters you passed.<br>
If you use a single optimizer it will iterate over everything, it will detect that gradients of non-used branch are none and that’s all. But it still uses time to iterate over all the parameters.</p>
<p>Using several optimizers will iterate over the which you actually know contain gradients. I don’t really know how many time you can save by doing this. Indeed, you would need to extend the logic of the forward pass.</p>
<p>So the normal way would be:</p>
<pre><code class="lang-auto">optimizer = OPTIM(model.parameters())
optimizer.zero_grad()
output=model(input)
loss = criterion(output,gt)
loss.backward()
optimizer.step()
</code></pre>
<p>It would become something like:</p>
<pre><code class="lang-auto">optimizer1 = OPTIM(branch1.parameters())
optimizer2 = OPTIM(branch2.parameters())
output=model(input)
loss = criterion(output,gt)
if x has PA view but not L view:
    optimizer1.zero_grad()
    loss.backward()
    optimizer1.step()
elif x has PA view but not L view:
    optimizer2.zero_grad()
    loss.backward()
    optimizer2.step()
</code></pre>
<p>But i don’t know how many time would you save.<br>
If you find it confusing just don’t do it.</p>