Repository: pytorch
Issue: Add a BlendLoss class in torch.nn and a blend_loss function in torch.nn.functional to allow blending of a list of loss functions
Labels: module: nn,module: loss,triaged,enhancement

## 🚀 Feature
<!-- A clear and concise description of the feature proposal -->

Add a `BlendLoss` class in `torch.nn` and a `blend_loss` function in `torch.nn.functional` to allow _out-of-the-box_ blending of a list of loss functions.

## Motivation

<!-- Please outline the motivation for the proposal. Is your feature request related to a problem? e.g., I'm always frustrated when [...]. If this is related to another GitHub issue, please link here too -->
Often one may require to blend multiple loss-functions to create a custom loss-function. While this could be done using standard mathematical operations, a way to blend various loss functions inside another dedicated class or function could come handy.

## Pitch

<!-- A clear and concise description of what you want to happen. -->

Simply put, the `BlendLoss` will do this:

> List of `loss_functions` + List of `weights` --> `BlendLoss` --> `blended_loss` as result.

![image](https://user-images.githubusercontent.com/10201242/120191554-b09ed880-c1df-11eb-9bff-44f92868e8d8.png)

<details>
<summary>
code for flowchart
</summary>
<p>

Test it on: https://mermaid-js.github.io/mermaid-live-editor

```mermaid
graph LR
    A1[List of loss_functions] -->|input| C1(nn.BlendLoss: nn.Module)
    B1[List of weights] -->|input| C1(nn.BlendLoss: nn.Module)
    C1 ==> CL([fa:fa-spinner a.forward])

    A2[List of loss_functions] -->|input| C2(nn.functional.blend_loss: function)
    B2[List of weights] -->|input| C2(nn.functional.blend_loss: function)
    C2 ==> CL([fa:fa-spinner blend losses])

    CL ===> |output| O[blend loss tensor] 

    subgraph BlendLoss: class
    A1
    B1
    C1
    end

    subgraph blend_loss: function
    A2
    B2
    C2
    end
```

</p>
</details>


## Suggested solution :fire: :star: 

<!-- A clear and concise description of any alternative solutions or features you've considered, if any. -->

Please find below a draft of the suggested preliminary code-solution. This only shows the possible implementation for the `class BlendLoss`. If this is accepted (or after some modifications), I will share/work the changes for the functional form of BlendLoss: `torch.nn.functional.blend_loss`. 

<details>
<summary>
<b>Expand to see code-solution </b> :zap:
</summary>
<p>

### **BlendLoss** class

```python
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import List, Union, Callable

class BlendLoss(nn.Module):
    """A loss function as a ``blend`` of a list of loss functions.
    """

    def __init__(self, 
                 loss_functions: List[Union[torch.nn.Module, Callable]],
                 weights: List[Union[float, int]] = None,                   
                 activation: Union[torch.nn.Module, Callable] = None
                 ):
        """Implementation of Blend loss for image segmentation task.
        
        Args:
            loss_functions: List of loss functions either as  
                    ``torch.nn.Module`` or ``torch.nn.functional.<loss>``.
            weights (optional): List of weights for the loss-functions. If `None` 
                    is provided, or, if the length of ``weights`` is different 
                    from that of ``loss_functions``, then equal weights are 
                    used for all loss-functions. (Default: None)
            activation (optional): An activation layer. Should be set to `None`
                    if your model contains a sigmoid or equivalent activation 
                    layer at the end of the model for prediction. (Default: None)

        """
        super(BlendLoss, self).__init__()
        self.loss_functions = loss_functions
        self.weights = weights
        self.activation = activation # example: torch.nn.functional.sigmoid
        self._update_loss_func_weights()


    def _update_loss_func_weights(self):
        """ Updates loss-function weights, when no weights are provided or when 
             ``len(weights)`` is not the same as ``len(loss_functions)``.
        """
        if self.loss_functions:            
            num_loss_funcs = len(self.loss_functions)
            if self.weights:
                # convert to tensors
                self.weights = torch.tensor(self.weights)
                # The number of weights must match the 
                # number loss-functions. If it does not, 
                # by default all weights will be set to 
                # ``1. / num_loss_funcs``.
                if len(self.weights) != num_loss_funcs:
                    self.weights = torch.ones(num_loss_funcs) 
                # normalize the weights
                self.weights = self.weights / self.weights.sum()
            else:
                self.weights = torch.ones(num_loss_funcs) / num_loss_funcs            
        else:            
            raise ValueError("No list of loss functions provided"
                " (each loss-function is either a Callable or torch.nn.Module).")


    def forward(self, inputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:

        assert_err_msg = "The inputs and targets cannot have different shapes." \
                         + "\n\t inputs.shape = {} \n\t targets.shape = {}"
        assert inputs.shape == targets.shape, assert_err_msg.format(
                                                inputs.shape, targets.shape)

        # The activation could be used when the model used 
        # does not have a final activation layer (e.g. sigmoid)
        # for prediction.
        if self.activation:
            inputs = self.activation(inputs)

        # Flatten label and prediction tensors
        inputs = inputs.view(-1)
        targets = targets.view(-1)

        # Apply loss-function blending
        blend = 0.
        for w, f in zip(self.weights, self.loss_functions):
            blend += w * f(inputs, targets)

        return blend
```

### Test `BlendLoss`

```python
# create test data
a = torch.arange(10.)
b = a + 2.
print(f'a: {a}')
print(f'a: {b}')

# define list of loss-functions
loss_functions = [nn.MSELoss(), nn.L1Loss()]
```

Test without any weights-specification. :star: 

```python
# without weights-specification
loss_functions = [nn.MSELoss(), nn.L1Loss()]
criterion = BlendLoss(loss_functions=loss_functions)
print(f'weights (normalized): {criterion.weights}')
loss = criterion.forward(inputs=a, targets=b)
print(f'loss: {loss: .2f}')
assert loss == 0.5 * nn.MSELoss()(a, b) + 0.5 * nn.L1Loss()(a, b)

## Output:
# weights (normalized): tensor([0.5000, 0.5000])
# loss: 3.00
```

Test with weights-specification. :star: 

```python
# with weights-specification
weights = [0.3, 0.7]
criterion = BlendLoss(loss_functions=loss_functions, weights=weights)
print(f'weights (normalized): {criterion.weights}')
loss = criterion.forward(inputs=a, targets=b)
print(f'loss: {loss: .2f}')
assert loss == (weights[0] * nn.MSELoss()(a, b) + 
                weights[1] * nn.L1Loss()(a, b)) / sum(weights)

## Output:
# weights (normalized): tensor([0.3000, 0.7000])
# loss:  2.60
```

</p>
</details>


cc @albanD @mruberry @jbschlosser