Add new Losses to VISSL¶
VISSL allows adding new losses easily. Follow the steps below to add a new loss:
Step1: Create a new loss
my_new_loss
invissl/losses/my_new_loss.py
following the template
import pprint
from classy_vision.losses import ClassyLoss, register_loss
@register_loss("my_new_loss")
class MyNewLoss(ClassyLoss):
"""
Add documentation for what the loss does
Config params:
document what parameters should be expected for the loss in the defaults.yaml
and how to set those params
"""
def __init__(self, loss_config: AttrDict, device: str = "gpu"):
super(MyNewLoss, self).__init__()
self.loss_config = loss_config
# implement what the init function should do
...
@classmethod
def from_config(cls, loss_config: AttrDict):
"""
Instantiates MyNewLoss from configuration.
Args:
loss_config: configuration for the loss
Returns:
MyNewLoss instance.
"""
return cls(loss_config)
def __repr__(self):
# implement what information about loss params should be
# printed by print(loss). This is helpful for debugging
repr_dict = {"name": self._get_name(), ....}
return pprint.pformat(repr_dict, indent=2)
def forward(self, output, target):
# implement how the loss should be calculated. The output should be
# torch.Tensor or List[torch.Tensor] and target should be torch.Tensor
...
...
return loss
Step2: Register the loss and loss params with VISSL Configuration. Add the params that the loss takes in VISSL defaults.yaml as follows:
LOSS:
my_new_loss:
param1: value1
param2: value2
...
Step3: Loss is ready to use. Simply set the configuration param
LOSS.name=my_new_loss