ml4gw 0.2.0__py3-none-any.whl → 0.4.0__py3-none-any.whl

ml4gw/nn/autoencoder/convolutional.py

@@ -0,0 +1,156 @@
+from collections.abc import Callable, Sequence
+from typing import Optional
+
+import torch
+
+from ml4gw.nn.autoencoder.base import Autoencoder
+from ml4gw.nn.autoencoder.skip_connection import SkipConnection
+from ml4gw.nn.autoencoder.utils import match_size
+
+Module = Callable[[...], torch.nn.Module]
+
+
+class ConvBlock(Autoencoder):
+    def __init__(
+        self,
+        in_channels: int,
+        encode_channels: int,
+        kernel_size: int,
+        stride: int = 1,
+        groups: int = 1,
+        activation: torch.nn.Module = torch.nn.ReLU(),
+        norm: Module = torch.nn.BatchNorm1d,
+        decode_channels: Optional[int] = None,
+        output_activation: Optional[torch.nn.Module] = None,
+        skip_connection: Optional[SkipConnection] = None,
+    ) -> None:
+        super().__init__(skip_connection=None)
+
+        self.kernel_size = kernel_size
+        self.padding = int((kernel_size - 1) // 2)
+        self.stride = stride
+
+        out_channels = encode_channels * groups
+        self.encode_layer = torch.nn.Conv1d(
+            in_channels,
+            out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=self.padding,
+            bias=False,
+            groups=groups,
+        )
+
+        decode_channels = decode_channels or in_channels
+        in_channels = encode_channels * groups
+        if skip_connection is not None:
+            in_channels = skip_connection.get_out_channels(in_channels)
+        self.decode_layer = torch.nn.ConvTranspose1d(
+            in_channels,
+            decode_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=self.padding,
+            bias=False,
+            groups=groups,
+        )
+
+        self.activation = activation
+        if output_activation is not None:
+            self.output_activation = output_activation
+        else:
+            self.output_activation = activation
+
+        self.encode_norm = norm(out_channels)
+        self.decode_norm = norm(decode_channels)
+
+    def encode(self, X):
+        X = self.encode_layer(X)
+        X = self.encode_norm(X)
+        return self.activation(X)
+
+    def decode(self, X):
+        X = self.decode_layer(X)
+        X = self.decode_norm(X)
+        return self.output_activation(X)
+
+
+class ConvolutionalAutoencoder(Autoencoder):
+    """
+    Build a stack of convolutional autoencoder layer
+    blocks. The output of each decoder layer will
+    match the shape of the input to its corresponding
+    encoder layer, except for the last decoder which
+    can have an arbitrary number of channels specified
+    by `decode_channels`.
+
+    All layers also share the same `activation` except
+    for the last decoder layer, which can have an
+    arbitrary `output_activation`.
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        encode_channels: Sequence[int],
+        kernel_size: int,
+        stride: int = 1,
+        groups: int = 1,
+        activation: torch.nn.Module = torch.nn.ReLU(),
+        output_activation: Optional[torch.nn.Module] = None,
+        norm: Module = torch.nn.BatchNorm1d,
+        decode_channels: Optional[int] = None,
+        skip_connection: Optional[SkipConnection] = None,
+    ) -> None:
+        # TODO: how to do this dynamically? Maybe the base
+        # architecture looks for overlapping arguments between
+        # this and the skip connection class and then provides them?
+        # if skip_connection is not None:
+        #     skip_connection = skip_connection(groups)
+        super().__init__(skip_connection=skip_connection)
+
+        output_activation = output_activation or activation
+        for i, channels in enumerate(encode_channels):
+            # All intermediate layers should decode to
+            # the same number of channels. The last decoder
+            # should decode to whatever number of channels
+            # was specified, even if it's `None` (in which
+            # case it will just be in_channels anyway)
+            decode = in_channels if i else decode_channels
+
+            # don't have the middle layer skip to itself
+            # TODO: wait I don't think this makes sense.
+            # j = len(encode_channels) - 1 - i
+            # connect = skip_connection if j else None
+            connect = skip_connection
+
+            # all intermediate layers should use the same
+            # activation. Only the last decoder should have
+            # a potentially different activation
+            out_act = None if i else output_activation
+
+            block = ConvBlock(
+                in_channels,
+                channels,
+                kernel_size,
+                stride,
+                groups,
+                activation=activation,
+                norm=norm,
+                decode_channels=decode,
+                skip_connection=connect,
+                output_activation=out_act,
+            )
+            self.blocks.append(block)
+            in_channels = channels * groups
+
+    def decode(self, *X, states=None, input_size: Optional[int] = None):
+        X = super().decode(*X, states=states)
+        if input_size is not None:
+            return match_size(X, input_size)
+        return X
+
+    def forward(self, X):
+        input_size = X.size(-1)
+        X = super().forward(X)
+        return match_size(X, input_size)

ml4gw/nn/autoencoder/skip_connection.py

@@ -0,0 +1,46 @@
+import torch
+
+from ml4gw.nn.autoencoder.utils import match_size
+
+
+class SkipConnection(torch.nn.Module):
+    def forward(self, X: torch.Tensor, state: torch.Tensor):
+        return match_size(X, state.size(-1))
+
+    def get_out_channels(self, in_channels):
+        return in_channels
+
+
+class AddSkipConnect(SkipConnection):
+    def forward(self, X, state):
+        X = super().forward(X, state)
+        return X + state
+
+
+class ConcatSkipConnect(SkipConnection):
+    def __init__(self, groups: int = 1):
+        super().__init__()
+        self.groups = groups
+
+    def get_out_channels(self, in_channels):
+        return 2 * in_channels
+
+    def forward(self, X, state):
+        X = super().forward(X, state)
+        if self.groups == 1:
+            return torch.cat([X, state], dim=1)
+
+        num_channels = X.size(1)
+        rem = num_channels % self.groups
+        if rem:
+            raise ValueError(
+                "Number of channels in input tensor {} cannot "
+                "be divided evenly into {} groups".format(
+                    num_channels, self.groups
+                )
+            )
+
+        X = torch.split(X, self.groups, dim=1)
+        state = torch.split(state, self.groups, dim=1)
+        frags = [i for j in zip(X, state) for i in j]
+        return torch.cat(frags, dim=1)

ml4gw/nn/autoencoder/utils.py

@@ -0,0 +1,14 @@
+import torch
+
+
+def match_size(X: torch.Tensor, target_size: int):
+    diff = target_size - X.size(-1)
+    left = int(diff // 2)
+    right = diff - left
+
+    if diff > 0:
+        return torch.nn.functional.pad(X, (left, right))
+    elif diff < 0:
+        right = -right or None
+        return X[:, :, -left:right]
+    return X

ml4gw/nn/norm.py

@@ -0,0 +1,97 @@
+from typing import Callable, Optional
+
+import torch
+
+NormLayer = Callable[[int], torch.nn.Module]
+
+
+class GroupNorm1D(torch.nn.Module):
+    """
+    Custom implementation of GroupNorm which is faster than the
+    out-of-the-box PyTorch version at inference time.
+    """
+
+    def __init__(
+        self,
+        num_channels: int,
+        num_groups: Optional[int] = None,
+        eps: float = 1e-5,
+    ):
+        super().__init__()
+        num_groups = num_groups or num_channels
+        if num_channels % num_groups:
+            raise ValueError("num_groups must be a factor of num_channels")
+
+        self.num_channels = num_channels
+        self.num_groups = num_groups
+        self.channels_per_group = self.num_channels // self.num_groups
+        self.eps = eps
+
+        shape = (self.num_channels, 1)
+        self.weight = torch.nn.Parameter(torch.ones(shape))
+        self.bias = torch.nn.Parameter(torch.zeros(shape))
+
+    def forward(self, x):
+        keepdims = self.num_groups == self.num_channels
+
+        # compute group variance via the E[x**2] - E**2[x] trick
+        mean = x.mean(-1, keepdims=keepdims)
+        sq_mean = (x**2).mean(-1, keepdims=keepdims)
+
+        # if we have groups, do some reshape magic
+        # to calculate group level stats then
+        # reshape back to full channel dimension
+        if self.num_groups != self.num_channels:
+            mean = torch.stack([mean, sq_mean], dim=1)
+            mean = mean.reshape(
+                -1, 2, self.num_groups, self.channels_per_group
+            )
+            mean = mean.mean(-1, keepdims=True)
+            mean = mean.expand(-1, -1, -1, self.channels_per_group)
+            mean = mean.reshape(-1, 2, self.num_channels, 1)
+            mean, sq_mean = mean[:, 0], mean[:, 1]
+
+        # roll the mean and variance into the
+        # weight and bias so that we have to do
+        # fewer computations along the full time axis
+        std = (sq_mean - mean**2 + self.eps) ** 0.5
+        scale = self.weight / std
+        shift = self.bias - scale * mean
+        return shift + x * scale
+
+
+class GroupNorm1DGetter:
+    """
+    Utility for making a NormLayer Callable that maps from
+    an integer number of channels to a torch Module. Useful
+    for command-line parameterization with jsonargparse.
+    """
+
+    def __init__(self, groups: Optional[int] = None) -> None:
+        self.groups = groups
+
+    def __call__(self, num_channels: int) -> torch.nn.Module:
+        if self.groups is None:
+            num_groups = None
+        else:
+            num_groups = min(num_channels, self.groups)
+        return GroupNorm1D(num_channels, num_groups)
+
+
+# TODO generalize faster 1dDGroupNorm to 2D
+class GroupNorm2DGetter:
+    """
+    Utility for making a NormLayer Callable that maps from
+    an integer number of channels to a torch Module. Useful
+    for command-line parameterization with jsonargparse.
+    """
+
+    def __init__(self, groups: Optional[int] = None) -> None:
+        self.groups = groups
+
+    def __call__(self, num_channels: int) -> torch.nn.Module:
+        if self.groups is None:
+            num_groups = num_channels
+        else:
+            num_groups = min(num_channels, self.groups)
+        return torch.nn.GroupNorm(num_groups, num_channels)

ml4gw/nn/resnet/__init__.py

@@ -0,0 +1,2 @@
+from .resnet_1d import ResNet1D
+from .resnet_2d import ResNet2D