Merge pull request #125 from sp-nitech/mdct

Add learnable option to MDCT

Author: takenori-y

README.md

@@ -2,7 +2,7 @@
 
 *diffsptk* is a differentiable version of [SPTK](https://github.com/sp-nitech/SPTK) based on the PyTorch framework.
 
-[![Manual](https://img.shields.io/badge/docs-stable-blue.svg)](https://sp-nitech.github.io/diffsptk/3.0.1/)
+[![Manual](https://img.shields.io/badge/docs-stable-blue.svg)](https://sp-nitech.github.io/diffsptk/3.0.2/)
 [![Downloads](https://static.pepy.tech/badge/diffsptk)](https://pepy.tech/project/diffsptk)
 [![ClickPy](https://img.shields.io/badge/downloads-clickpy-yellow.svg)](https://clickpy.clickhouse.com/dashboard/diffsptk)
 [![Python Version](https://img.shields.io/pypi/pyversions/diffsptk.svg)](https://pypi.python.org/pypi/diffsptk)
@@ -20,7 +20,7 @@
 
 ## Documentation
 
-- See [this page](https://sp-nitech.github.io/diffsptk/3.0.1/) for the reference manual.
+- See [this page](https://sp-nitech.github.io/diffsptk/3.0.2/) for the reference manual.
 - Our [paper](https://www.isca-speech.org/archive/ssw_2023/yoshimura23_ssw.html) is available on the ISCA Archive.
 
 ## Installation

diffsptk/modules/dfs.py

@@ -53,7 +53,7 @@ def __init__(
     ) -> None:
         super().__init__()
 
-        _, _, tensors = self._precompute(*get_values(locals(), end=-2))
+        _, _, tensors = self._precompute(*get_values(locals()))
         if learnable and b is not None:
             self.b = nn.Parameter(tensors[0])
         else:

diffsptk/modules/drc.py

@@ -74,7 +74,7 @@ def __init__(
     ) -> None:
         super().__init__()
 
-        self.values, _, tensors = self._precompute(*get_values(locals(), end=-2))
+        self.values, _, tensors = self._precompute(*get_values(locals()))
         if learnable:
             self.params = nn.Parameter(tensors[0])
         else:

diffsptk/modules/imdct.py

@@ -20,7 +20,7 @@
 from ..typing import Callable, Precomputed
 from ..utils.private import check_size, get_layer, get_values
 from .base import BaseFunctionalModule
-from .mdct import ModifiedDiscreteTransform
+from .mdct import LEARNABLES, ModifiedDiscreteCosineTransform, ModifiedDiscreteTransform
 from .unframe import Unframe
 from .window import Window
 
@@ -36,12 +36,22 @@ class InverseModifiedDiscreteCosineTransform(BaseFunctionalModule):
     window : ['sine', 'vorbis', 'kbd', 'rectangular']
         The window type.
 
+    learnable : bool or list[str]
+        Indicates whether the parameters are learnable. If a boolean, it specifies
+        whether all parameters are learnable. If a list, it contains the keys of the
+        learnable parameters, which can only be "basis" and "window".
+
     """
 
-    def __init__(self, frame_length: int, window: str = "sine") -> None:
+    def __init__(
+        self,
+        frame_length: int,
+        window: str = "sine",
+        learnable: bool | list[str] = False,
+    ) -> None:
         super().__init__()
 
-        self.values, layers, _ = self._precompute(*get_values(locals()))
+        self.values, layers, _ = self._precompute(*get_values(locals(), full=True))
         self.layers = nn.ModuleList(layers)
 
     def forward(self, y: torch.Tensor, out_length: int | None = None) -> torch.Tensor:
@@ -89,23 +99,33 @@ def _takes_input_size() -> bool:
         return False
 
     @staticmethod
-    def _check() -> None:
-        pass
+    def _check(*args, **kwargs) -> None:
+        ModifiedDiscreteCosineTransform._check(*args, **kwargs)
 
     @staticmethod
     def _precompute(
-        frame_length: int, window: str, transform: str = "cosine", module: bool = True
+        frame_length: int,
+        window: str,
+        learnable: bool = False,
+        transform: str = "cosine",
+        module: bool = True,
     ) -> Precomputed:
-        InverseModifiedDiscreteCosineTransform._check()
+        InverseModifiedDiscreteCosineTransform._check(learnable)
         frame_period = frame_length // 2
 
+        if learnable is True:
+            learnable = LEARNABLES
+        elif learnable is False:
+            learnable = ()
+
         imdt = get_layer(
             module,
             InverseModifiedDiscreteTransform,
             dict(
                 length=frame_length,
                 window=window,
                 transform=transform,
+                learnable="basis" in learnable,
             ),
         )
         window_ = get_layer(
@@ -116,6 +136,7 @@ def _precompute(
                 out_length=None,
                 window=window,
                 norm="none",
+                learnable="window" in learnable,
             ),
         )
         unframe = get_layer(
@@ -157,15 +178,27 @@ class InverseModifiedDiscreteTransform(BaseFunctionalModule):
     transform : ['cosine', 'sine']
         The transform type.
 
+    learnable : bool
+        Whether to make the DCT matrix learnable.
+
     """
 
-    def __init__(self, length: int, window: str, transform: str = "cosine") -> None:
+    def __init__(
+        self,
+        length: int,
+        window: str,
+        transform: str = "cosine",
+        learnable: bool = False,
+    ) -> None:
         super().__init__()
 
         self.in_dim = length // 2
 
         _, _, tensors = self._precompute(*get_values(locals()))
-        self.register_buffer("W", tensors[0])
+        if learnable:
+            self.W = nn.Parameter(tensors[0])
+        else:
+            self.register_buffer("W", tensors[0])
 
     def forward(self, y: torch.Tensor) -> torch.Tensor:
         """Apply inverse MDCT/MDST to the input.
@@ -182,7 +215,7 @@ def forward(self, y: torch.Tensor) -> torch.Tensor:
 
         """
         check_size(y.size(-1), self.in_dim, "dimension of input")
-        return self._forward(y, **self._buffers)
+        return self._forward(y, **self._buffers, **self._parameters)
 
     @staticmethod
     def _func(y: torch.Tensor, *args, **kwargs) -> torch.Tensor:

diffsptk/modules/imdst.py

@@ -34,12 +34,22 @@ class InverseModifiedDiscreteSineTransform(BaseFunctionalModule):
     window : ['sine', 'vorbis', 'kbd', 'rectangular']
         The window type.
 
+    learnable : bool or list[str]
+        Indicates whether the parameters are learnable. If a boolean, it specifies
+        whether all parameters are learnable. If a list, it contains the keys of the
+        learnable parameters, which can only be "basis" and "window".
+
     """
 
-    def __init__(self, frame_length: int, window: str = "sine"):
+    def __init__(
+        self,
+        frame_length: int,
+        window: str = "sine",
+        learnable: bool | list[str] = False,
+    ):
         super().__init__()
 
-        self.values, layers, _ = self._precompute(*get_values(locals()))
+        self.values, layers, _ = self._precompute(*get_values(locals(), full=True))
         self.layers = nn.ModuleList(layers)
 
     def forward(self, y: torch.Tensor, out_length: int | None = None) -> torch.Tensor:
@@ -86,8 +96,8 @@ def _check(*args, **kwargs) -> None:
         raise NotImplementedError
 
     @staticmethod
-    def _precompute(frame_length: int, window: str) -> Precomputed:
-        return IMDCT._precompute(frame_length, window, transform="sine")
+    def _precompute(*args, **kwargs) -> Precomputed:
+        return IMDCT._precompute(*args, **kwargs, transform="sine")
 
     @staticmethod
     def _forward(*args, **kwargs) -> torch.Tensor:

diffsptk/modules/mdct.py

@@ -24,6 +24,8 @@
 from .frame import Frame
 from .window import Window
 
+LEARNABLES = ("basis", "window")
+
 
 class ModifiedDiscreteCosineTransform(BaseFunctionalModule):
     """This module is a simple cascade of framing, windowing, and modified DCT.
@@ -36,12 +38,22 @@ class ModifiedDiscreteCosineTransform(BaseFunctionalModule):
     window : ['sine', 'vorbis', 'kbd', 'rectangular']
         The window type.
 
+    learnable : bool or list[str]
+        Indicates whether the parameters are learnable. If a boolean, it specifies
+        whether all parameters are learnable. If a list, it contains the keys of the
+        learnable parameters, which can only be "basis" and "window".
+
     """
 
-    def __init__(self, frame_length: int, window: str = "sine") -> None:
+    def __init__(
+        self,
+        frame_length: int,
+        window: str = "sine",
+        learnable: bool | list[str] = False,
+    ) -> None:
         super().__init__()
 
-        self.values, layers, _ = self._precompute(*get_values(locals()))
+        self.values, layers, _ = self._precompute(*get_values(locals(), full=True))
         self.layers = nn.ModuleList(layers)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -84,16 +96,29 @@ def _takes_input_size() -> bool:
         return False
 
     @staticmethod
-    def _check() -> None:
-        pass
+    def _check(learnable: bool | list[str]) -> None:
+        if isinstance(learnable, (tuple, list)):
+            if any(x not in LEARNABLES for x in learnable):
+                raise ValueError("An unsupported key is found in learnable.")
+        elif not isinstance(learnable, bool):
+            raise ValueError("learnable must be boolean or list.")
 
     @staticmethod
     def _precompute(
-        frame_length: int, window: str, transform: str = "cosine", module: bool = True
+        frame_length: int,
+        window: str,
+        learnable: bool | list[str] = False,
+        transform: str = "cosine",
+        module: bool = True,
     ) -> Precomputed:
-        ModifiedDiscreteCosineTransform._check()
+        ModifiedDiscreteCosineTransform._check(learnable)
         frame_period = frame_length // 2
 
+        if learnable is True:
+            learnable = LEARNABLES
+        elif learnable is False:
+            learnable = ()
+
         frame = get_layer(
             module,
             Frame,
@@ -110,6 +135,7 @@ def _precompute(
                 out_length=None,
                 window=window,
                 norm="none",
+                learnable="window" in learnable,
             ),
         )
         mdt = get_layer(
@@ -119,6 +145,7 @@ def _precompute(
                 length=frame_length,
                 window=window,
                 transform=transform,
+                learnable="basis" in learnable,
             ),
         )
         return (frame_period,), (frame, window_, mdt), None
@@ -150,15 +177,27 @@ class ModifiedDiscreteTransform(BaseFunctionalModule):
     transform : ['cosine', 'sine']
         The transform type.
 
+    learnable : bool
+        Whether to make the DCT matrix learnable.
+
     """
 
-    def __init__(self, length: int, window: str, transform: str = "cosine") -> None:
+    def __init__(
+        self,
+        length: int,
+        window: str,
+        transform: str = "cosine",
+        learnable: bool = False,
+    ) -> None:
         super().__init__()
 
         self.in_dim = length
 
         _, _, tensors = self._precompute(*get_values(locals()))
-        self.register_buffer("W", tensors[0])
+        if learnable:
+            self.W = nn.Parameter(tensors[0])
+        else:
+            self.register_buffer("W", tensors[0])
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         """Apply MDCT/MDST to the input.
@@ -175,7 +214,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
         """
         check_size(x.size(-1), self.in_dim, "dimension of input")
-        return self._forward(x, **self._buffers)
+        return self._forward(x, **self._buffers, **self._parameters)
 
     @staticmethod
     def _func(x: torch.Tensor, *args, **kwargs) -> torch.Tensor:

diffsptk/modules/mdst.py

@@ -34,12 +34,22 @@ class ModifiedDiscreteSineTransform(BaseFunctionalModule):
     window : ['sine', 'vorbis', 'kbd', 'rectangular']
         The window type.
 
+    learnable : bool or list[str]
+        Indicates whether the parameters are learnable. If a boolean, it specifies
+        whether all parameters are learnable. If a list, it contains the keys of the
+        learnable parameters, which can only be "basis" and "window".
+
     """
 
-    def __init__(self, frame_length: int, window: str = "sine") -> None:
+    def __init__(
+        self,
+        frame_length: int,
+        window: str = "sine",
+        learnable: bool | list[str] = False,
+    ) -> None:
         super().__init__()
 
-        self.values, layers, _ = self._precompute(*get_values(locals()))
+        self.values, layers, _ = self._precompute(*get_values(locals(), full=True))
         self.layers = nn.ModuleList(layers)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -83,8 +93,8 @@ def _check(*args, **kwargs) -> None:
         raise NotImplementedError
 
     @staticmethod
-    def _precompute(frame_length: int, window: str) -> Precomputed:
-        return MDCT._precompute(frame_length, window, transform="sine")
+    def _precompute(*args, **kwargs) -> Precomputed:
+        return MDCT._precompute(*args, **kwargs, transform="sine")
 
     @staticmethod
     def _forward(*args, **kwargs) -> torch.Tensor:

diffsptk/modules/window.py

@@ -16,6 +16,7 @@
 
 import torch
 import torch.nn.functional as F
+from torch import nn
 
 from ..typing import Precomputed
 from ..utils.private import check_size, get_values, to
@@ -42,6 +43,9 @@ class Window(BaseFunctionalModule):
     norm : ['none', 'power', 'magnitude']
         The normalization type of the window.
 
+    learnable : bool
+        Whether to make the window learnable.
+
     """
 
     def __init__(
@@ -51,13 +55,17 @@ def __init__(
         *,
         window: str | int = "blackman",
         norm: str | int = "power",
+        learnable: bool = False,
     ) -> None:
         super().__init__()
 
         self.in_dim = in_length
 
         self.values, _, tensors = self._precompute(*get_values(locals()))
-        self.register_buffer("window", tensors[0])
+        if learnable:
+            self.window = nn.Parameter(tensors[0])
+        else:
+            self.register_buffer("window", tensors[0])
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         """Apply a window function to the given waveform.
@@ -82,7 +90,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
         """
         check_size(x.size(-1), self.in_dim, "input length")
-        return self._forward(x, *self.values, **self._buffers)
+        return self._forward(x, *self.values, **self._buffers, **self._parameters)
 
     @staticmethod
     def _func(x: torch.Tensor, *args, **kwargs) -> torch.Tensor: