torch-utils 1.0.0__py3-none-any.whl

torch_utils/__init__.py

@@ -0,0 +1,4 @@
+import importlib.metadata
+
+
+__version__ = importlib.metadata.version('torch_utils')

torch_utils/_layer.py

@@ -0,0 +1,145 @@
+from __future__ import annotations
+
+from typing import ClassVar
+
+import torch
+from torch import nn
+
+
+__all__ = ['Conv1d', 'Conv2d', 'Conv3d', 'Deconv2d', 'Layer', 'Linear']
+
+
+class Layer(nn.Module):
+    class Inherit:
+        __slots__ = ()
+
+    # Default layer arguments
+    ACTIVATION: type[nn.Module] | None = nn.LeakyReLU
+    ACTIVATION_KWARGS: ClassVar[dict] = {'negative_slope': 0.1}
+
+    USE_BATCH_NORM = True
+    BATCH_NORM_TRAINING = True
+    BATCH_NORM_MOMENTUM = 0.01
+
+    def __init__(self, activation: type[nn.Module] | nn.Module | type[Layer.Inherit] | None):
+        super().__init__()
+        # Preload default
+        if activation == Layer.Inherit:
+            activation = Layer.ACTIVATION
+        self.activation: nn.Module | None = (
+            activation(**Layer.ACTIVATION_KWARGS) if isinstance(activation, type) else activation)  # ty: ignore[invalid-assignment]
+        self.batch_norm: nn.modules.batchnorm._BatchNorm | None = None
+
+    def forward(self, input_data: torch.Tensor) -> torch.Tensor:
+        output = input_data
+        if self.activation:
+            output = self.activation(output)
+        if self.batch_norm:
+            output = self.batch_norm(output)
+        return output
+
+    @staticmethod
+    def add_weight_decay(module: nn.Module, weight_decay: float, exclude=()):
+        decay = []
+        no_decay = []
+        for name, param in module.named_parameters():
+            if not param.requires_grad:
+                continue
+            if len(param.shape) == 1 or name.endswith('.bias') or name in exclude:
+                no_decay.append(param)
+            else:
+                decay.append(param)
+        return [
+            {'params': no_decay, 'weight_decay': 0.0},
+            {'params': decay, 'weight_decay': weight_decay}]
+
+
+class Linear(Layer):
+    def __init__(self, in_channels: int, out_channels: int,
+                 activation: nn.Module | type[Layer.Inherit] | None = Layer.Inherit,
+                 use_batch_norm: bool | None = None, **kwargs):
+        super().__init__(activation)
+
+        use_batch_norm = Layer.USE_BATCH_NORM if use_batch_norm is None else use_batch_norm
+        self.linear = nn.Linear(in_channels, out_channels, bias=not use_batch_norm, **kwargs)
+        self.batch_norm = nn.BatchNorm1d(
+            out_channels,
+            momentum=Layer.BATCH_NORM_MOMENTUM,
+            track_running_stats=Layer.BATCH_NORM_TRAINING) if use_batch_norm else None
+
+    def forward(self, input_data: torch.Tensor) -> torch.Tensor:
+        return super().forward(self.linear(input_data))
+
+
+class Conv1d(Layer):
+    def __init__(self, in_channels: int, out_channels: int, kernel_size: int = 3,
+                 stride: int | tuple[int] = 1, activation: nn.Module | type[Layer.Inherit] | None = Layer.Inherit,
+                 use_batch_norm: bool | None = None, **kwargs):
+        super().__init__(activation)
+
+        use_batch_norm = Layer.USE_BATCH_NORM if use_batch_norm is None else use_batch_norm
+        self.conv = nn.Conv1d(in_channels, out_channels, kernel_size, stride=stride,
+                              bias=use_batch_norm, **kwargs)
+        self.batch_norm = nn.BatchNorm1d(
+            out_channels,
+            momentum=Layer.BATCH_NORM_MOMENTUM,
+            track_running_stats=Layer.BATCH_NORM_TRAINING) if use_batch_norm else None
+
+    def forward(self, input_data: torch.Tensor) -> torch.Tensor:
+        return super().forward(self.conv(input_data))
+
+
+class Conv2d(Layer):
+    def __init__(self, in_channels: int, out_channels: int, kernel_size: int | tuple[int, int] = 3,
+                 stride: int | tuple[int, int] = 1, activation: nn.Module | type[Layer.Inherit] | None = Layer.Inherit,
+                 use_batch_norm: bool | None = None, **kwargs):
+        super().__init__(activation)
+
+        use_batch_norm = Layer.USE_BATCH_NORM if use_batch_norm is None else use_batch_norm
+        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride=stride,
+                              bias=not use_batch_norm, **kwargs)
+        self.batch_norm = nn.BatchNorm2d(
+            out_channels,
+            momentum=Layer.BATCH_NORM_MOMENTUM,
+            track_running_stats=Layer.BATCH_NORM_TRAINING) if use_batch_norm else None
+
+    def forward(self, input_data: torch.Tensor) -> torch.Tensor:
+        return super().forward(self.conv(input_data))
+
+
+class Conv3d(Layer):
+    def __init__(self, in_channels: int, out_channels: int, kernel_size: int | tuple[int, int, int] = 3,
+                 stride: int | tuple[int, int, int] = 1,
+                 activation: nn.Module | type[Layer.Inherit] | None = Layer.Inherit, use_batch_norm: bool | None = None,
+                 **kwargs):
+        super().__init__(activation)
+
+        use_batch_norm = Layer.USE_BATCH_NORM if use_batch_norm is None else use_batch_norm
+        self.conv = nn.Conv3d(in_channels, out_channels, kernel_size, stride=stride,
+                              bias=use_batch_norm, **kwargs)
+        self.batch_norm = nn.BatchNorm3d(
+            out_channels,
+            momentum=Layer.BATCH_NORM_MOMENTUM,
+            track_running_stats=Layer.BATCH_NORM_TRAINING) if use_batch_norm else None
+
+    def forward(self, input_data: torch.Tensor) -> torch.Tensor:
+        return super().forward(self.conv(input_data))
+
+
+class Deconv2d(Layer):
+    def __init__(self, in_channels: int, out_channels: int, kernel_size: int = 3,
+                 stride: int | tuple[int, int] = 1, activation: nn.Module | type[Layer.Inherit] | None = Layer.Inherit,
+                 use_batch_norm: bool | None = None, **kwargs):
+        super().__init__(activation)
+
+        use_batch_norm = Layer.USE_BATCH_NORM if use_batch_norm is None else use_batch_norm
+        self.deconv = nn.ConvTranspose2d(
+            in_channels, out_channels, kernel_size, stride=stride,
+            bias=not use_batch_norm, **kwargs)
+        self.batch_norm = nn.BatchNorm2d(
+            out_channels,
+            momentum=Layer.BATCH_NORM_MOMENTUM,
+            track_running_stats=Layer.BATCH_NORM_TRAINING) if use_batch_norm else None
+
+    def forward(self, input_data: torch.Tensor) -> torch.Tensor:
+        return super().forward(self.deconv(input_data))

torch_utils/_residual.py

@@ -0,0 +1,55 @@
+import torch
+from torch import nn
+
+from ._layer import Conv2d, Layer
+
+
+__all__ = ['ResBlock', 'ResBottleneck']
+
+
+class ResBlock(Layer):
+    def __init__(self, in_channels: int, out_channels: int = -1, kernel_size: int = 3, padding: int = 1,
+                 stride: int | tuple[int, int] = 1, activation: nn.Module | type[Layer.Inherit] | None = Layer.Inherit,
+                 batch_norm=None, **kwargs):
+        super().__init__(activation)
+        self.batch_norm = None
+        if out_channels == -1:
+            out_channels = in_channels
+
+        self.seq = nn.Sequential(
+            Conv2d(in_channels, in_channels, kernel_size=kernel_size, stride=stride, padding=padding, **kwargs),
+            Conv2d(in_channels, out_channels, kernel_size=3, padding=1,
+                   activation=None, batch_norm=batch_norm))
+        self.residual = Conv2d(in_channels, out_channels, kernel_size=1, stride=stride, activation=None) if (
+            out_channels != in_channels or stride != 1) else None
+
+    def forward(self, input_data: torch.Tensor) -> torch.Tensor:
+        if self.residual is not None:
+            return super().forward(self.residual(input_data) + self.seq(input_data))
+        return super().forward(input_data + self.seq(input_data))
+
+
+class ResBottleneck(Layer):
+    def __init__(self, in_channels: int, out_channels: int = -1, bottleneck_channels: int = -1, kernel_size: int = 3,
+                 stride: int | tuple[int, int] = 1, padding=1,
+                 activation: nn.Module | type[Layer.Inherit] | None = Layer.Inherit, batch_norm=None, **kwargs):
+        super().__init__(activation)
+        self.batch_norm = None
+        if out_channels == -1:
+            out_channels = in_channels
+        if bottleneck_channels == -1:
+            bottleneck_channels = in_channels // 4
+
+        self.seq = nn.Sequential(
+            Conv2d(in_channels, bottleneck_channels, kernel_size=1),
+            Conv2d(bottleneck_channels, bottleneck_channels, kernel_size=kernel_size,
+                   stride=stride, padding=padding, **kwargs),
+            Conv2d(bottleneck_channels, out_channels, kernel_size=1,
+                   activation=None, batch_norm=batch_norm))
+        self.residual = Conv2d(in_channels, out_channels, kernel_size=1, stride=stride, activation=None) if (
+            out_channels != in_channels or stride != 1) else None
+
+    def forward(self, input_data: torch.Tensor) -> torch.Tensor:
+        if self.residual is not None:
+            return super().forward(self.residual(input_data) + self.seq(input_data))
+        return super().forward(input_data + self.seq(input_data))

torch_utils/batch/__init__.py

@@ -0,0 +1,6 @@
+from .batch_generator import BatchGenerator
+from .ipc_batch_generator import IPCBatchGenerator
+from .sequence_batch_generator import SequenceBatchGenerator
+
+
+__all__ = ['BatchGenerator', 'IPCBatchGenerator', 'SequenceBatchGenerator']