Repository: facebookresearch/NeuralCompression (language: python)

Issue:
Rate distortion loss 
## Enhancement

A rate distortion loss 

## Motivation

Rate distortion losses are used by a variety of models 
Rate distortion loss 
## Enhancement

A rate distortion loss 

## Motivation

Rate distortion losses are used by a variety of models 


Discussion / hints:

Fix patch:
diff --git a/docs/source/layers.rst b/docs/source/layers.rst
index 58f1d408..961773c5 100644
--- a/docs/source/layers.rst
+++ b/docs/source/layers.rst
@@ -45,6 +45,13 @@ NonNegativeParameterization
   :members: __getitem__
   :special-members:
 
+RateMSEDistortionLoss
+~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: RateMSEDistortionLoss
+  :members: __getitem__
+  :special-members:
+
 SimplifiedGDN
 ~~~~~~~~~~~~~
 
diff --git a/neuralcompression/layers/__init__.py b/neuralcompression/layers/__init__.py
index 96dfc1ca..ec141610 100644
--- a/neuralcompression/layers/__init__.py
+++ b/neuralcompression/layers/__init__.py
@@ -9,5 +9,6 @@
 from ._hyper_analysis_transformation_2d import HyperAnalysisTransformation2D
 from ._hyper_synthesis_transformation_2d import HyperSynthesisTransformation2D
 from ._non_negative_parameterization import NonNegativeParameterization
+from ._rate_mse_distortion_loss import RateMSEDistortionLoss
 from ._synthesis_transformation_2d import SynthesisTransformation2D
 from .gdn import SimplifiedGDN, SimplifiedInverseGDN
diff --git a/neuralcompression/layers/_rate_mse_distortion_loss.py b/neuralcompression/layers/_rate_mse_distortion_loss.py
new file mode 100644
index 00000000..9ac7fab7
--- /dev/null
+++ b/neuralcompression/layers/_rate_mse_distortion_loss.py
@@ -0,0 +1,78 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import math
+from typing import List, NamedTuple
+
+import torch
+from torch import Tensor
+from torch.nn import Module, MSELoss
+
+
+class RateDistortionLoss(NamedTuple):
+    distortion: float
+    rate: float
+    rate_distortion: float
+
+
+class RateMSEDistortionLoss(Module):
+    """Rate-distortion loss.
+
+    The rate-distortion loss is the minimum transmission bit-rate for a
+    required quality. It can be obtained without consideration of a specific
+    coding method. Rate is expressed in bits per pixel (BPP) of the original,
+    ``x``, distortion is expressed as the mean squared error (MSE) between the
+    original, ``x``, and the target, ``x_hat``.
+
+    Args:
+        trade_off: rate-distortion trade-off. `trade_off = 1` is the solution
+            where the `(rate, distortion)` pair minimizes `rate + distortion`.
+            Increasing `trade_off` will penalize the distortion term so more
+            bits are spent.
+        maximum: dynamic range of the input (i.e. the difference between the
+            maximum the and minimum permitted values).
+    """
+
+    def __init__(self, trade_off: float = 1e-2, maximum: int = 255):
+        super(RateMSEDistortionLoss, self).__init__()
+
+        self._maximum = maximum
+
+        self._trade_off = trade_off
+
+        self._mse = MSELoss()
+
+    def forward(
+        self,
+        x_hat: Tensor,
+        probabilities: List[Tensor],
+        x: Tensor,
+    ) -> RateDistortionLoss:
+        """
+        Args:
+            x_hat: encoder output.
+            probabilities: reconstruction likelihoods.
+            x: encoder input.
+        """
+        n, _, h, w = x.size()
+
+        bpps = []
+
+        for probability in probabilities:
+            pixels = -math.log(2) * (n * h * w)
+
+            bpps += [float(torch.log(probability).sum() / pixels)]
+
+        rate = sum(bpps)
+
+        distortion = self._mse(x_hat, x)
+
+        rate_distortion = rate + distortion
+
+        return RateDistortionLoss(
+            rate,
+            distortion,
+            self._trade_off * self._maximum ** 2 * rate_distortion,
+        )


Test patch:
diff --git a/tests/layers/test_rate_mse_distortion_loss.py b/tests/layers/test_rate_mse_distortion_loss.py
new file mode 100644
index 00000000..56d4f03f
--- /dev/null
+++ b/tests/layers/test_rate_mse_distortion_loss.py
@@ -0,0 +1,28 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+
+from neuralcompression.layers import RateMSEDistortionLoss
+
+
+class TestRateMSEDistortionLoss:
+    def test_forward(self):
+        x_hat = torch.ones(1, 32, 32, 3)
+
+        scores = [
+            torch.ones((32,)),
+            torch.ones((32,)),
+        ]
+
+        x = torch.ones(1, 32, 32, 3)
+
+        rate_distortion_loss = RateMSEDistortionLoss()
+
+        rate_distortion_losses = rate_distortion_loss.forward(x_hat, scores, x)
+
+        assert rate_distortion_losses.rate == 0.0
+        assert rate_distortion_losses.distortion == 0.0
+        assert rate_distortion_losses.rate_distortion == 0.0