How to implement nn.Module.forward for both train and eval mode?

username: epignatelli
cooked: <p>I am trying to embed the loss calculation in the model itself, rather than attaching the module at every iteration.<br>
This requires the <code>forward</code> function to have another argument in the signature<br>
<code>forward(self, x, y)</code>, in order to have the info to calculate the loss.</p>
<p>What is the correct way to handle this for both <code>train</code> and <code>eval</code> mode?</p>
<p>Is defaulting <code>y=None</code> the recommended way of doing so?</p>
<p>I wasn’t able to find any documentation on it, any reference would be appreciated!</p>

username: Mazhar_Shaikh
cooked: <p>The “.train()” and “.eval()” mode are related to certain modules like dropout, batchnorms, etc which have a different functionality in the two modes. However, They do not change the signature of the function call of the model.<br>
To embed the loss calculation in the model itself, the defaulting to y=None should work. There is no recommended way to do this. You could follow the template in <a href="https://github.com/ipod825/recurrent-visual-attention/blob/master/model.py" rel="nofollow noopener">this repository</a>.</p>

username: epignatelli
cooked: <p>Thanks Mazhar_Shaikh, I chose that way, as it was the simplest implementation.</p>
<p>Just to give another example, the <code>torchvision.models.detection.Generalized_RCNN</code> model uses the same pattern, and complements it with a further check for <code>model.training</code> at line 44.</p>
<aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/pytorch/vision/blob/8635be94d1216f10fb8302da89233bd86445e449/torchvision/models/detection/generalized_rcnn.py#L31-L62" target="_blank" rel="nofollow noopener">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/vision/blob/8635be94d1216f10fb8302da89233bd86445e449/torchvision/models/detection/generalized_rcnn.py#L31-L62" target="_blank" rel="nofollow noopener">pytorch/vision/blob/8635be94d1216f10fb8302da89233bd86445e449/torchvision/models/detection/generalized_rcnn.py#L31-L62</a></h4>
<pre class="onebox"><code class="lang-py"><ol class="start lines" start="31" style="counter-reset: li-counter 30 ;">
<li>def forward(self, images, targets=None):</li>
<li>    """</li>
<li>    Arguments:</li>
<li>        images (list[Tensor]): images to be processed</li>
<li>        targets (list[Dict[Tensor]]): ground-truth boxes present in the image (optional)</li>
<li>
</li>
<li>    Returns:</li>
<li>        result (list[BoxList] or dict[Tensor]): the output from the model.</li>
<li>            During training, it returns a dict[Tensor] which contains the losses.</li>
<li>            During testing, it returns list[BoxList] contains additional fields</li>
<li>            like `scores`, `labels` and `mask` (for Mask R-CNN models).</li>
<li>
</li>
<li>    """</li>
<li>    if self.training and targets is None:</li>
<li>        raise ValueError("In training mode, targets should be passed")</li>
<li>    original_image_sizes = [img.shape[-2:] for img in images]</li>
<li>    images, targets = self.transform(images, targets)</li>
<li>    features = self.backbone(images.tensors)</li>
<li>    if isinstance(features, torch.Tensor):</li>
<li>        features = OrderedDict([(0, features)])</li>
</ol></code></pre>

  This file has been truncated. <a href="https://github.com/pytorch/vision/blob/8635be94d1216f10fb8302da89233bd86445e449/torchvision/models/detection/generalized_rcnn.py#L31-L62" target="_blank" rel="nofollow noopener">show original</a>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>