Repository: Project-MONAI/MONAI
File: monai/losses/barlow_twins.py (Python)

Code before:
# Copyright (c) MONAI Consortium
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#     http://www.apache.org/licenses/LICENSE-2.0
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations

from warnings import warn

import torch
from torch.nn.modules.loss import _Loss


class BarlowTwinsLoss(_Loss):
    """
    Compute the Barlow Twins loss defined in:

        Zbontar, Jure, et al. "Barlow Twins: Self-Supervised Learning via Redundancy Reduction" International
        conference on machine learning. PMLR, 2020. (http://proceedings.mlr.press/v139/zbontar21a/zbontar21a.pdf)

    Adapted from:
        https://github.com/facebookresearch/barlowtwins

    """

    def __init__(self, lambd: float = 5e-3, batch_size: int = -1) -> None:
        """
        Args:
            lamb: Can be any float to handle the informativeness and invariance trade-off. Ideally set to 5e-3.

        Raises:
            ValueError: When an input of dimension length > 2 is passed
            ValueError: When input and target are of different shapes
            ValueError: When batch size is less than or equal to 1

        """
        super().__init__()
        self.lambd = lambd

        if batch_size != -1:
            warn("batch_size is no longer required to be set. It will be estimated dynamically in the forward call")

Review comment (suggestion):
Thanks for the contribution! Could you please provide some explanation here in the docstring on what the loss is for, what it computes, what applications it would be for, example use case if there's something particular to demonstrate, etc.? It helps at a glance to know what this loss would be for without reading the paper. 

Code after:
# Copyright (c) MONAI Consortium
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#     http://www.apache.org/licenses/LICENSE-2.0
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations

import torch
from torch.nn.modules.loss import _Loss


class BarlowTwinsLoss(_Loss):
    """
    The Barlow Twins cost function takes the representations extracted by a neural network from two
    distorted views and seeks to make the cross-correlation matrix of the two representations tend
    towards identity. This encourages the neural network to learn similar representations with the least
    amount of redundancy. This cost function can be used in particular in multimodal learning to work on
    representations from two modalities. The most common use case is for unsupervised learning, where data
    augmentations are used to generate 2 distorted views of the same sample to force the encoder to
    extract useful features for downstream tasks.

    Zbontar, Jure, et al. "Barlow Twins: Self-Supervised Learning via Redundancy Reduction" International
    conference on machine learning. PMLR, 2020. (http://proceedings.mlr.press/v139/zbontar21a/zbontar21a.pdf)

    Adapted from:
        https://github.com/facebookresearch/barlowtwins

    """

    def __init__(self, lambd: float = 5e-3) -> None:
        """
        Args:
            lamb: Can be any float to handle the informativeness and invariance trade-off. Ideally set to 5e-3.

        Raises:
            ValueError: When an input of dimension length > 2 is passed
            ValueError: When input and target are of different shapes
            ValueError: When batch size is less than or equal to 1

        """
        super().__init__()