refactor: LatentCodec eliminate Optional[*]s

Author: YodaEmbedding

compressai/latent_codecs/base.py

@@ -38,45 +38,7 @@
 ]
 
 
-class _SetDefaultMixin:
-    """Convenience functions for initializing classes with defaults."""
-
-    def _setdefault(self, k, v, f):
-        """Initialize attribute ``k`` with value ``v`` or ``f()``."""
-        v = v or f()
-        setattr(self, k, v)
-
-    # TODO instead of save_direct, override load_state_dict() and state_dict()
-    def _set_group_defaults(self, group_key, group_dict, defaults, save_direct=False):
-        """Initialize attribute ``group_key`` with items from
-        ``group_dict``, using defaults for missing keys.
-        Ensures ``nn.Module`` attributes are properly registered.
-
-        Args:
-            - group_key:
-                Name of attribute.
-            - group_dict:
-                Dict of items to initialize ``group_key`` with.
-            - defaults:
-                Dict of defaults for items not in ``group_dict``.
-            - save_direct:
-                If ``True``, save items directly as attributes of ``self``.
-                If ``False``, save items in a ``nn.ModuleDict``.
-        """
-        group_dict = group_dict if group_dict is not None else {}
-        for k, f in defaults.items():
-            if k in group_dict:
-                continue
-            group_dict[k] = f()
-        if save_direct:
-            for k, v in group_dict.items():
-                setattr(self, k, v)
-        else:
-            group_dict = nn.ModuleDict(group_dict)
-        setattr(self, group_key, group_dict)
-
-
-class LatentCodec(nn.Module, _SetDefaultMixin):
+class LatentCodec(nn.Module):
     def forward(self, y: Tensor, *args, **kwargs) -> Dict[str, Any]:
         raise NotImplementedError
 

compressai/latent_codecs/channel_groups.py

@@ -28,7 +28,7 @@
 # ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
 from itertools import accumulate
-from typing import Any, Dict, List, Mapping, Optional, Tuple
+from typing import Any, Dict, List, Mapping, Tuple
 
 import torch
 import torch.nn as nn
@@ -70,8 +70,8 @@ class ChannelGroupsLatentCodec(LatentCodec):
 
     def __init__(
         self,
-        latent_codec: Optional[Mapping[str, LatentCodec]] = None,
-        channel_context: Optional[Mapping[str, nn.Module]] = None,
+        latent_codec: Mapping[str, LatentCodec],
+        channel_context: Mapping[str, nn.Module],
         *,
         groups: List[int],
         **kwargs,

compressai/latent_codecs/checkerboard.py

@@ -27,7 +27,7 @@
 # OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
 # ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
-from typing import Any, Dict, List, Mapping, Optional, Tuple
+from typing import Any, Dict, List, Mapping, Tuple
 
 import torch
 import torch.nn as nn
@@ -40,7 +40,6 @@
 from compressai.registry import register_module
 
 from .base import LatentCodec
-from .gaussian_conditional import GaussianConditionalLatentCodec
 
 __all__ = [
     "CheckerboardLatentCodec",
@@ -109,16 +108,11 @@ class CheckerboardLatentCodec(LatentCodec):
         □   empty
     """
 
-    latent_codec: Mapping[str, LatentCodec]
-
-    entropy_parameters: nn.Module
-    context_prediction: CheckerboardMaskedConv2d
-
     def __init__(
         self,
-        latent_codec: Optional[Mapping[str, LatentCodec]] = None,
-        entropy_parameters: Optional[nn.Module] = None,
-        context_prediction: Optional[nn.Module] = None,
+        latent_codec: Mapping[str, LatentCodec],
+        entropy_parameters: nn.Module,
+        context_prediction: CheckerboardMaskedConv2d,
         anchor_parity="even",
         forward_method="twopass",
         **kwargs,
@@ -128,16 +122,10 @@ def __init__(
         self.anchor_parity = anchor_parity
         self.non_anchor_parity = {"odd": "even", "even": "odd"}[anchor_parity]
         self.forward_method = forward_method
-        self.entropy_parameters = entropy_parameters or nn.Identity()
-        self.context_prediction = context_prediction or nn.Identity()
-        self._set_group_defaults(
-            "latent_codec",
-            latent_codec,
-            defaults={
-                "y": lambda: GaussianConditionalLatentCodec(quantizer="ste"),
-            },
-            save_direct=True,
-        )
+        self.entropy_parameters = entropy_parameters
+        self.context_prediction = context_prediction
+        self.y = latent_codec["y"]
+        self.latent_codec = latent_codec
 
     def __getitem__(self, key: str) -> LatentCodec:
         return self.latent_codec[key]

compressai/latent_codecs/gain/hyper.py

@@ -27,7 +27,7 @@
 # OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
 # ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Dict, List, Tuple
 
 import torch.nn as nn
 
@@ -66,22 +66,17 @@ class GainHyperLatentCodec(LatentCodec):
 
     """
 
-    entropy_bottleneck: EntropyBottleneck
-    h_a: nn.Module
-    h_s: nn.Module
-
     def __init__(
         self,
-        entropy_bottleneck: Optional[EntropyBottleneck] = None,
-        h_a: Optional[nn.Module] = None,
-        h_s: Optional[nn.Module] = None,
+        entropy_bottleneck: EntropyBottleneck,
+        h_a: nn.Module,
+        h_s: nn.Module,
         **kwargs,
     ):
         super().__init__()
-        assert entropy_bottleneck is not None
         self.entropy_bottleneck = entropy_bottleneck
-        self.h_a = h_a or nn.Identity()
-        self.h_s = h_s or nn.Identity()
+        self.h_a = h_a
+        self.h_s = h_s
 
     def forward(self, y: Tensor, gain: Tensor, gain_inv: Tensor) -> Dict[str, Any]:
         z = self.h_a(y)

compressai/latent_codecs/gain/hyperprior.py

@@ -27,15 +27,13 @@
 # OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
 # ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
-from typing import Any, Dict, List, Mapping, Optional, Tuple
+from typing import Any, Dict, List, Mapping, Tuple
 
 from torch import Tensor
 
 from compressai.registry import register_module
 
 from ..base import LatentCodec
-from ..gaussian_conditional import GaussianConditionalLatentCodec
-from .hyper import GainHyperLatentCodec
 
 __all__ = [
     "GainHyperpriorLatentCodec",
@@ -94,21 +92,11 @@ class GainHyperpriorLatentCodec(LatentCodec):
      - entropy bottleneck ``hyper`` (default) and autoregressive ``y``
     """
 
-    latent_codec: Mapping[str, LatentCodec]
-
-    def __init__(
-        self, latent_codec: Optional[Mapping[str, LatentCodec]] = None, **kwargs
-    ):
+    def __init__(self, latent_codec: Mapping[str, LatentCodec], **kwargs):
         super().__init__()
-        self._set_group_defaults(
-            "latent_codec",
-            latent_codec,
-            defaults={
-                "y": GaussianConditionalLatentCodec,
-                "hyper": GainHyperLatentCodec,
-            },
-            save_direct=True,
-        )
+        self.y = latent_codec["y"]
+        self.hyper = latent_codec["hyper"]
+        self.latent_codec = latent_codec
 
     def __getitem__(self, key: str) -> LatentCodec:
         return self.latent_codec[key]

compressai/latent_codecs/hyper.py

@@ -27,7 +27,7 @@
 # OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
 # ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Dict, List, Tuple
 
 import torch.nn as nn
 
@@ -65,23 +65,18 @@ class HyperLatentCodec(LatentCodec):
 
     """
 
-    entropy_bottleneck: EntropyBottleneck
-    h_a: nn.Module
-    h_s: nn.Module
-
     def __init__(
         self,
-        entropy_bottleneck: Optional[EntropyBottleneck] = None,
-        h_a: Optional[nn.Module] = None,
-        h_s: Optional[nn.Module] = None,
+        entropy_bottleneck: EntropyBottleneck,
+        h_a: nn.Module,
+        h_s: nn.Module,
         quantizer: str = "noise",
         **kwargs,
     ):
         super().__init__()
-        assert entropy_bottleneck is not None
         self.entropy_bottleneck = entropy_bottleneck
-        self.h_a = h_a or nn.Identity()
-        self.h_s = h_s or nn.Identity()
+        self.h_a = h_a
+        self.h_s = h_s
         self.quantizer = quantizer
 
     def forward(self, y: Tensor) -> Dict[str, Any]:

compressai/latent_codecs/hyperprior.py

@@ -27,15 +27,13 @@
 # OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
 # ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
-from typing import Any, Dict, List, Mapping, Optional, Tuple
+from typing import Any, Dict, List, Mapping, Tuple
 
 from torch import Tensor
 
 from compressai.registry import register_module
 
 from .base import LatentCodec
-from .gaussian_conditional import GaussianConditionalLatentCodec
-from .hyper import HyperLatentCodec
 
 __all__ = [
     "HyperpriorLatentCodec",
@@ -87,21 +85,11 @@ class HyperpriorLatentCodec(LatentCodec):
      - entropy bottleneck ``hyper`` (default) and autoregressive ``y``
     """
 
-    latent_codec: Mapping[str, LatentCodec]
-
-    def __init__(
-        self, latent_codec: Optional[Mapping[str, LatentCodec]] = None, **kwargs
-    ):
+    def __init__(self, latent_codec: Mapping[str, LatentCodec], **kwargs):
         super().__init__()
-        self._set_group_defaults(
-            "latent_codec",
-            latent_codec,
-            defaults={
-                "y": GaussianConditionalLatentCodec,
-                "hyper": HyperLatentCodec,
-            },
-            save_direct=True,
-        )
+        self.y = latent_codec["y"]
+        self.hyper = latent_codec["hyper"]
+        self.latent_codec = latent_codec
 
     def __getitem__(self, key: str) -> LatentCodec:
         return self.latent_codec[key]

compressai/latent_codecs/rasterscan.py

@@ -27,7 +27,7 @@
 # OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
 # ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
-from typing import Any, Callable, Dict, List, Optional, Tuple, TypeVar
+from typing import Any, Callable, Dict, List, Tuple, TypeVar
 
 import torch
 import torch.nn as nn
@@ -83,21 +83,17 @@ class RasterScanLatentCodec(LatentCodec):
 
     """
 
-    gaussian_conditional: GaussianConditional
-    entropy_parameters: nn.Module
-    context_prediction: MaskedConv2d
-
     def __init__(
         self,
-        gaussian_conditional: Optional[GaussianConditional] = None,
-        entropy_parameters: Optional[nn.Module] = None,
-        context_prediction: Optional[MaskedConv2d] = None,
+        gaussian_conditional: GaussianConditional,
+        entropy_parameters: nn.Module,
+        context_prediction: MaskedConv2d,
         **kwargs,
     ):
         super().__init__()
-        self.gaussian_conditional = gaussian_conditional or GaussianConditional(None)
-        self.entropy_parameters = entropy_parameters or nn.Identity()
-        self.context_prediction = context_prediction or MaskedConv2d()
+        self.gaussian_conditional = gaussian_conditional
+        self.entropy_parameters = entropy_parameters
+        self.context_prediction = context_prediction
         self.kernel_size = _to_single(self.context_prediction.kernel_size)
         self.padding = (self.kernel_size - 1) // 2
 

docs/source/latent_codecs.rst

@@ -231,6 +231,7 @@ Using :py:class:`~compressai.models.base.SimpleVAECompressionModel`, some Google
                     ),
                     # Encode y using autoregression in raster-scan order:
                     "y": RasterScanLatentCodec(
+                        gaussian_conditional=GaussianConditional(None),
                         entropy_parameters=nn.Sequential(...),
                         context_prediction=MaskedConv2d(
                             M, M * 2, kernel_size=5, padding=2, stride=1
@@ -249,7 +250,7 @@ Latent codecs should inherit from the abstract base class :py:class:`~LatentCode
 
 .. code-block:: python
 
-    class LatentCodec(nn.Module, _SetDefaultMixin):
+    class LatentCodec(nn.Module):
         def forward(self, y: Tensor, *args, **kwargs) -> Dict[str, Any]:
             raise NotImplementedError