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

epignatelli · August 10, 2019, 8:15am

I am trying to embed the loss calculation in the model itself, rather than attaching the module at every iteration.
This requires the forward function to have another argument in the signature
forward(self, x, y), in order to have the info to calculate the loss.

What is the correct way to handle this for both train and eval mode?

Is defaulting y=None the recommended way of doing so?

I wasn’t able to find any documentation on it, any reference would be appreciated!

Mazhar_Shaikh · August 10, 2019, 10:05am

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.
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 this repository.

epignatelli · August 10, 2019, 12:46pm

Thanks Mazhar_Shaikh, I chose that way, as it was the simplest implementation.

Just to give another example, the torchvision.models.detection.Generalized_RCNN model uses the same pattern, and complements it with a further check for model.training at line 44.

github.com

pytorch/vision/blob/8635be94d1216f10fb8302da89233bd86445e449/torchvision/models/detection/generalized_rcnn.py#L31-L62


def forward(self, images, targets=None):
    """
    Arguments:
        images (list[Tensor]): images to be processed
        targets (list[Dict[Tensor]]): ground-truth boxes present in the image (optional)


    Returns:
        result (list[BoxList] or dict[Tensor]): the output from the model.
            During training, it returns a dict[Tensor] which contains the losses.
            During testing, it returns list[BoxList] contains additional fields
            like `scores`, `labels` and `mask` (for Mask R-CNN models).


    """
    if self.training and targets is None:
        raise ValueError("In training mode, targets should be passed")
    original_image_sizes = [img.shape[-2:] for img in images]
    images, targets = self.transform(images, targets)
    features = self.backbone(images.tensors)
    if isinstance(features, torch.Tensor):
        features = OrderedDict([(0, features)])

This file has been truncated. show original