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

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

ml4gw/nn/resnet/resnet_1d.py

@@ -0,0 +1,413 @@
+"""
+In large part lifted from
+https://github.com/pytorch/vision/blob/main/torchvision/models/resnet.py
+but with 1d convolutions and arbitrary kernel sizes, and a
+default norm layer that makes more sense for most GW applications
+where training-time statistics are entirely arbitrary due to
+simulations.
+"""
+
+from typing import Callable, List, Literal, Optional
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+from ml4gw.nn.norm import GroupNorm1DGetter, NormLayer
+
+
+def convN(
+    in_planes: int,
+    out_planes: int,
+    kernel_size: int = 3,
+    stride: int = 1,
+    groups: int = 1,
+    dilation: int = 1,
+) -> nn.Conv1d:
+    """1d convolution with padding"""
+    if not kernel_size % 2:
+        raise ValueError("Can't use even sized kernels")
+
+    return nn.Conv1d(
+        in_planes,
+        out_planes,
+        kernel_size=kernel_size,
+        stride=stride,
+        padding=dilation * int(kernel_size // 2),
+        groups=groups,
+        bias=False,
+        dilation=dilation,
+    )
+
+
+def conv1(in_planes: int, out_planes: int, stride: int = 1) -> nn.Conv1d:
+    """Kernel-size 1 convolution"""
+    return nn.Conv1d(
+        in_planes, out_planes, kernel_size=1, stride=stride, bias=False
+    )
+
+
+class BasicBlock(nn.Module):
+    """Defines the structure of the blocks used to build the ResNet"""
+
+    expansion: int = 1
+
+    def __init__(
+        self,
+        inplanes: int,
+        planes: int,
+        kernel_size: int = 3,
+        stride: int = 1,
+        downsample: Optional[nn.Module] = None,
+        groups: int = 1,
+        base_width: int = 64,
+        dilation: int = 1,
+        norm_layer: Optional[Callable[..., nn.Module]] = None,
+    ) -> None:
+
+        super().__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm1d
+        if groups != 1 or base_width != 64:
+            raise ValueError(
+                "BasicBlock only supports groups=1 and base_width=64"
+            )
+        if dilation > 1:
+            raise NotImplementedError(
+                "Dilation > 1 not supported in BasicBlock"
+            )
+
+        # Both self.conv1 and self.downsample layers
+        # downsample the input when stride != 1
+        self.conv1 = convN(inplanes, planes, kernel_size, stride)
+        self.bn1 = norm_layer(planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = convN(planes, planes, kernel_size)
+        self.bn2 = norm_layer(planes)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x: Tensor) -> Tensor:
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+
+class Bottleneck(nn.Module):
+    """
+    Bottleneck blocks implement one extra convolution
+    compared to basic blocks. In this layers, the `planes`
+    parameter is generally meant to _downsize_ the number
+    of feature maps first, which then get expanded out to
+    `planes * Bottleneck.expansion` feature maps at the
+    output of the layer.
+    """
+
+    expansion: int = 4
+
+    def __init__(
+        self,
+        inplanes: int,
+        planes: int,
+        kernel_size: int = 3,
+        stride: int = 1,
+        downsample: Optional[nn.Module] = None,
+        groups: int = 1,
+        base_width: int = 64,
+        dilation: int = 1,
+        norm_layer: Optional[NormLayer] = None,
+    ) -> None:
+        super().__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm1d
+
+        width = int(planes * (base_width / 64.0)) * groups
+
+        # conv1 does no downsampling, just reduces the number of
+        # feature maps from inplanes to width (where width == planes)
+        # if groups == 1 and base_width == 64
+        self.conv1 = convN(inplanes, width, kernel_size)
+        self.bn1 = norm_layer(width)
+
+        # conv2 keeps the same number of feature maps,
+        # but downsamples along the time axis if stride
+        # or dilation > 1
+        self.conv2 = convN(width, width, kernel_size, stride, groups, dilation)
+        self.bn2 = norm_layer(width)
+
+        # conv3 expands the feature maps back out to planes * expansion
+        self.conv3 = conv1(width, planes * self.expansion)
+        self.bn3 = norm_layer(planes * self.expansion)
+
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x: Tensor) -> Tensor:
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+
+class ResNet1D(nn.Module):
+    """1D ResNet architecture
+
+    Simple extension of ResNet to 1D convolutions with
+    arbitrary kernel sizes to support the longer timeseries
+    used in BBH detection.
+
+    Args:
+        in_channels:
+            The number of channels in input tensor.
+        layers:
+            A list representing the number of residual
+            blocks to include in each "layer" of the
+            network. Total layers (e.g. 50 in ResNet50)
+            is `2 + sum(layers) * factor`, where factor
+            is `2` for vanilla `ResNet` and `3` for
+            `BottleneckResNet`.
+        kernel_size:
+            The size of the convolutional kernel to
+            use in all residual layers. _NOT_ the size
+            of the input kernel to the network, which
+            is determined at run-time.
+        zero_init_residual:
+            Flag indicating whether to initialize the
+            weights of the batch-norm layer in each block
+            to 0 so that residuals are initialized as
+            identities. Can improve training results.
+        groups:
+            Number of convolutional groups to use in all
+            layers. Grouped convolutions induce local
+            connections between feature maps at subsequent
+            layers rather than global. Generally won't
+            need this to be >1, and wil raise an error if
+            >1 when using vanilla `ResNet`.
+        width_per_group:
+            Base width of each of the feature map groups,
+            which is scaled up by the typical expansion
+            factor at each layer of the network. Meaningless
+            for vanilla `ResNet`.
+        stride_type:
+            Whether to achieve downsampling on the time axis
+            by strided or dilated convolutions for each layer.
+            If left as `None`, strided convolutions will be
+            used at each layer. Otherwise, `stride_type` should
+            be one element shorter than `layers` and indicate either
+            `stride` or `dilation` for each layer after the first.
+    """
+
+    block = BasicBlock
+
+    def __init__(
+        self,
+        in_channels: int,
+        layers: List[int],
+        classes: int,
+        kernel_size: int = 3,
+        zero_init_residual: bool = False,
+        groups: int = 1,
+        width_per_group: int = 64,
+        stride_type: Optional[List[Literal["stride", "dilation"]]] = None,
+        norm_layer: Optional[NormLayer] = None,
+    ) -> None:
+        super().__init__()
+
+        self.inplanes = 64
+        self.dilation = 1
+
+        # default to using InstanceNorm if no
+        # norm layer is provided explicitly
+        self._norm_layer = norm_layer or GroupNorm1DGetter()
+
+        # TODO: should we support passing a single string
+        # for simplicity here?
+        if stride_type is None:
+            # each element in the tuple indicates if we should replace
+            # the stride with a dilated convolution instead
+            stride_type = ["stride"] * (len(layers) - 1)
+        if len(stride_type) != (len(layers) - 1):
+            raise ValueError(
+                "'stride_type' should be None or a "
+                "{}-element tuple, got {}".format(len(layers) - 1, stride_type)
+            )
+
+        self.groups = groups
+        self.base_width = width_per_group
+
+        # start with a basic conv-bn-relu-maxpool block
+        # to reduce the dimensionality before the heavy
+        # lifting starts
+        self.conv1 = nn.Conv1d(
+            in_channels,
+            self.inplanes,
+            kernel_size=7,
+            stride=2,
+            padding=3,
+            bias=False,
+        )
+        self.bn1 = self._norm_layer(self.inplanes)
+        self.relu = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool1d(kernel_size=3, stride=2, padding=1)
+
+        # now create layers of residual blocks where each
+        # layer uses the same number of feature maps for
+        # all its blocks (some power of 2 times 64).
+        # Don't downsample along the time axis in the first
+        # layer, but downsample in all the rest (either by
+        # striding or dilating depending on the stride_type
+        # argument)
+        residual_layers = [self._make_layer(64, layers[0], kernel_size)]
+        it = zip(layers[1:], stride_type)
+        for i, (num_blocks, stride) in enumerate(it):
+            block_size = 64 * 2 ** (i + 1)
+            layer = self._make_layer(
+                block_size,
+                num_blocks,
+                kernel_size,
+                stride=2,
+                stride_type=stride,
+            )
+            residual_layers.append(layer)
+        self.residual_layers = nn.ModuleList(residual_layers)
+
+        # Average pool over each feature map to create a
+        # single value for each feature map that we'll use
+        # in the fully connected head
+        self.avgpool = nn.AdaptiveAvgPool1d(1)
+
+        # use a fully connected layer to map from the
+        # feature maps to the binary output that we need
+        self.fc = nn.Linear(block_size * self.block.expansion, classes)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv1d):
+                nn.init.kaiming_normal_(
+                    m.weight, mode="fan_out", nonlinearity="relu"
+                )
+            elif isinstance(m, (nn.BatchNorm1d, nn.GroupNorm)):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+        # Zero-initialize the last BN in each residual branch,
+        # so that the residual branch starts with zeros,
+        # and each residual block behaves like an identity.
+        # This improves the model by 0.2~0.3% according to
+        # https://arxiv.org/abs/1706.02677
+        if zero_init_residual:
+            for m in self.modules():
+                if isinstance(m, Bottleneck):
+                    nn.init.constant_(m.bn3.weight, 0)
+                elif isinstance(m, BasicBlock):
+                    nn.init.constant_(m.bn2.weight, 0)
+
+    def _make_layer(
+        self,
+        planes: int,
+        blocks: int,
+        kernel_size: int = 3,
+        stride: int = 1,
+        stride_type: Literal["stride", "dilation"] = "stride",
+    ) -> nn.Sequential:
+        block = self.block
+        norm_layer = self._norm_layer
+        downsample = None
+        previous_dilation = self.dilation
+
+        if stride_type == "dilation":
+            self.dilation *= stride
+            stride = 1
+        elif stride_type != "stride":
+            raise ValueError("Unknown stride type {stride}")
+
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                conv1(self.inplanes, planes * block.expansion, stride),
+                norm_layer(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(
+            block(
+                self.inplanes,
+                planes,
+                kernel_size,
+                stride,
+                downsample,
+                self.groups,
+                self.base_width,
+                previous_dilation,
+                norm_layer,
+            )
+        )
+        self.inplanes = planes * block.expansion
+        for _ in range(1, blocks):
+            layers.append(
+                block(
+                    self.inplanes,
+                    planes,
+                    kernel_size,
+                    groups=self.groups,
+                    base_width=self.base_width,
+                    dilation=self.dilation,
+                    norm_layer=norm_layer,
+                )
+            )
+
+        return nn.Sequential(*layers)
+
+    def _forward_impl(self, x: Tensor) -> Tensor:
+        # See note [TorchScript super()]
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+
+        for layer in self.residual_layers:
+            x = layer(x)
+
+        x = self.avgpool(x)
+        x = torch.flatten(x, 1)
+        x = self.fc(x)
+
+        return x
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self._forward_impl(x)
+
+
+# TODO: implement as arg of ResNet instead?
+class BottleneckResNet1D(ResNet1D):
+    """A version of ResNet that uses bottleneck blocks"""
+
+    block = Bottleneck

ml4gw/nn/resnet/resnet_2d.py

@@ -0,0 +1,413 @@
+"""
+In large part lifted from
+https://github.com/pytorch/vision/blob/main/torchvision/models/resnet.py
+but with arbitrary kernel sizes
+"""
+
+from typing import Callable, List, Literal, Optional
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+from ml4gw.nn.norm import GroupNorm2DGetter, NormLayer
+
+
+def convN(
+    in_planes: int,
+    out_planes: int,
+    kernel_size: int = 3,
+    stride: int = 1,
+    groups: int = 1,
+    dilation: int = 1,
+) -> nn.Conv2d:
+    """2d convolution with padding"""
+    if not kernel_size % 2:
+        raise ValueError("Can't use even sized kernels")
+
+    return nn.Conv2d(
+        in_planes,
+        out_planes,
+        kernel_size=kernel_size,
+        stride=stride,
+        padding=dilation * int(kernel_size // 2),
+        groups=groups,
+        bias=False,
+        dilation=dilation,
+    )
+
+
+def conv1(in_planes: int, out_planes: int, stride: int = 1) -> nn.Conv2d:
+    """Kernel-size 1 convolution"""
+    return nn.Conv2d(
+        in_planes, out_planes, kernel_size=1, stride=stride, bias=False
+    )
+
+
+class BasicBlock(nn.Module):
+    """Defines the structure of the blocks used to build the ResNet"""
+
+    expansion: int = 1
+
+    def __init__(
+        self,
+        inplanes: int,
+        planes: int,
+        kernel_size: int = 3,
+        stride: int = 1,
+        downsample: Optional[nn.Module] = None,
+        groups: int = 1,
+        base_width: int = 64,
+        dilation: int = 1,
+        norm_layer: Optional[Callable[..., nn.Module]] = None,
+    ) -> None:
+
+        super().__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        if groups != 1 or base_width != 64:
+            raise ValueError(
+                "BasicBlock only supports groups=1 and base_width=64"
+            )
+        if dilation > 1:
+            raise NotImplementedError(
+                "Dilation > 1 not supported in BasicBlock"
+            )
+
+        # Both self.conv1 and self.downsample layers
+        # downsample the input when stride != 1
+        self.conv1 = convN(inplanes, planes, kernel_size, stride)
+        self.bn1 = norm_layer(planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = convN(planes, planes, kernel_size)
+        self.bn2 = norm_layer(planes)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x: Tensor) -> Tensor:
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+
+class Bottleneck(nn.Module):
+    """
+    Bottleneck blocks implement one extra convolution
+    compared to basic blocks. In this layers, the `planes`
+    parameter is generally meant to _downsize_ the number
+    of feature maps first, which then get expanded out to
+    `planes * Bottleneck.expansion` feature maps at the
+    output of the layer.
+    """
+
+    expansion: int = 4
+
+    def __init__(
+        self,
+        inplanes: int,
+        planes: int,
+        kernel_size: int = 3,
+        stride: int = 1,
+        downsample: Optional[nn.Module] = None,
+        groups: int = 1,
+        base_width: int = 64,
+        dilation: int = 1,
+        norm_layer: Optional[Callable[..., nn.Module]] = None,
+    ) -> None:
+        super().__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+
+        width = int(planes * (base_width / 64.0)) * groups
+
+        # conv1 does no downsampling, just reduces the number of
+        # feature maps from inplanes to width (where width == planes)
+        # if groups == 1 and base_width == 64
+        self.conv1 = convN(inplanes, width, kernel_size)
+        self.bn1 = norm_layer(width)
+
+        # conv2 keeps the same number of feature maps,
+        # but downsamples along the time axis if stride
+        # or dilation > 1
+        self.conv2 = convN(width, width, kernel_size, stride, groups, dilation)
+        self.bn2 = norm_layer(width)
+
+        # conv3 expands the feature maps back out to planes * expansion
+        self.conv3 = conv1(width, planes * self.expansion)
+        self.bn3 = norm_layer(planes * self.expansion)
+
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x: Tensor) -> Tensor:
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+
+class ResNet2D(nn.Module):
+    """2D ResNet architecture
+
+    Simple extension of ResNet with arbitrary kernel sizes
+    to support the longer timeseries used in BBH detection.
+
+    Args:
+        in_channels:
+            The number of channels in input tensor.
+        layers:
+            A list representing the number of residual
+            blocks to include in each "layer" of the
+            network. Total layers (e.g. 50 in ResNet50)
+            is `2 + sum(layers) * factor`, where factor
+            is `2` for vanilla `ResNet` and `3` for
+            `BottleneckResNet`.
+        kernel_size:
+            The size of the convolutional kernel to
+            use in all residual layers. _NOT_ the size
+            of the input kernel to the network, which
+            is determined at run-time.
+        zero_init_residual:
+            Flag indicating whether to initialize the
+            weights of the batch-norm layer in each block
+            to 0 so that residuals are initialized as
+            identities. Can improve training results.
+        groups:
+            Number of convolutional groups to use in all
+            layers. Grouped convolutions induce local
+            connections between feature maps at subsequent
+            layers rather than global. Generally won't
+            need this to be >1, and wil raise an error if
+            >1 when using vanilla `ResNet`.
+        width_per_group:
+            Base width of each of the feature map groups,
+            which is scaled up by the typical expansion
+            factor at each layer of the network. Meaningless
+            for vanilla `ResNet`.
+        stride_type:
+            Whether to achieve downsampling on the time axis
+            by strided or dilated convolutions for each layer.
+            If left as `None`, strided convolutions will be
+            used at each layer. Otherwise, `stride_type` should
+            be one element shorter than `layers` and indicate either
+            `stride` or `dilation` for each layer after the first.
+        norm_groups:
+            The number of groups to use in GroupNorm layers
+            throughout the model. If left as `-1`, the number
+            of groups will be equal to the number of channels,
+            making this equilavent to LayerNorm
+    """
+
+    block = BasicBlock
+
+    def __init__(
+        self,
+        in_channels: int,
+        layers: List[int],
+        classes: int,
+        kernel_size: int = 3,
+        zero_init_residual: bool = False,
+        groups: int = 1,
+        width_per_group: int = 64,
+        stride_type: Optional[List[Literal["stride", "dilation"]]] = None,
+        norm_layer: Optional[NormLayer] = None,
+    ) -> None:
+        super().__init__()
+        # default to using InstanceNorm if no
+        # norm layer is provided explicitly
+        self._norm_layer = norm_layer or GroupNorm2DGetter()
+
+        self.inplanes = 64
+        self.dilation = 1
+
+        # TODO: should we support passing a single string
+        # for simplicity here?
+        if stride_type is None:
+            # each element in the tuple indicates if we should replace
+            # the stride with a dilated convolution instead
+            stride_type = ["stride"] * (len(layers) - 1)
+        if len(stride_type) != (len(layers) - 1):
+            raise ValueError(
+                "'stride_type' should be None or a "
+                "{}-element tuple, got {}".format(len(layers) - 1, stride_type)
+            )
+
+        self.groups = groups
+        self.base_width = width_per_group
+
+        # start with a basic conv-bn-relu-maxpool block
+        # to reduce the dimensionality before the heavy
+        # lifting starts
+        self.conv1 = nn.Conv2d(
+            in_channels,
+            self.inplanes,
+            kernel_size=7,
+            stride=2,
+            padding=3,
+            bias=False,
+        )
+        self.bn1 = self._norm_layer(self.inplanes)
+        self.relu = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+
+        # now create layers of residual blocks where each
+        # layer uses the same number of feature maps for
+        # all its blocks (some power of 2 times 64).
+        # Don't downsample along the time axis in the first
+        # layer, but downsample in all the rest (either by
+        # striding or dilating depending on the stride_type
+        # argument)
+        residual_layers = [self._make_layer(64, layers[0], kernel_size)]
+        it = zip(layers[1:], stride_type)
+        for i, (num_blocks, stride) in enumerate(it):
+            block_size = 64 * 2 ** (i + 1)
+            layer = self._make_layer(
+                block_size,
+                num_blocks,
+                kernel_size,
+                stride=2,
+                stride_type=stride,
+            )
+            residual_layers.append(layer)
+        self.residual_layers = nn.ModuleList(residual_layers)
+
+        # Average pool over each feature map to create a
+        # single value for each feature map that we'll use
+        # in the fully connected head
+        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+
+        # use a fully connected layer to map from the
+        # feature maps to the binary output that we need
+        self.fc = nn.Linear(block_size * self.block.expansion, classes)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(
+                    m.weight, mode="fan_out", nonlinearity="relu"
+                )
+            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+        # Zero-initialize the last BN in each residual branch,
+        # so that the residual branch starts with zeros,
+        # and each residual block behaves like an identity.
+        # This improves the model by 0.2~0.3% according to
+        # https://arxiv.org/abs/1706.02677
+        if zero_init_residual:
+            for m in self.modules():
+                if isinstance(m, Bottleneck):
+                    nn.init.constant_(m.bn3.weight, 0)
+                elif isinstance(m, BasicBlock):
+                    nn.init.constant_(m.bn2.weight, 0)
+
+    def _make_layer(
+        self,
+        planes: int,
+        blocks: int,
+        kernel_size: int = 3,
+        stride: int = 1,
+        stride_type: Literal["stride", "dilation"] = "stride",
+    ) -> nn.Sequential:
+        block = self.block
+        norm_layer = self._norm_layer
+        downsample = None
+        previous_dilation = self.dilation
+
+        if stride_type == "dilation":
+            self.dilation *= stride
+            stride = 1
+        elif stride_type != "stride":
+            raise ValueError("Unknown stride type {stride}")
+
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                conv1(self.inplanes, planes * block.expansion, stride),
+                norm_layer(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(
+            block(
+                self.inplanes,
+                planes,
+                kernel_size,
+                stride,
+                downsample,
+                self.groups,
+                self.base_width,
+                previous_dilation,
+                norm_layer,
+            )
+        )
+        self.inplanes = planes * block.expansion
+        for _ in range(1, blocks):
+            layers.append(
+                block(
+                    self.inplanes,
+                    planes,
+                    kernel_size,
+                    groups=self.groups,
+                    base_width=self.base_width,
+                    dilation=self.dilation,
+                    norm_layer=norm_layer,
+                )
+            )
+
+        return nn.Sequential(*layers)
+
+    def _forward_impl(self, x: Tensor) -> Tensor:
+        # See note [TorchScript super()]
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+
+        for layer in self.residual_layers:
+            x = layer(x)
+
+        x = self.avgpool(x)
+        x = torch.flatten(x, 1)
+        x = self.fc(x)
+
+        return x
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self._forward_impl(x)
+
+
+# TODO: implement as arg of ResNet instead?
+class BottleneckResNet2D(ResNet2D):
+    """A version of ResNet that uses bottleneck blocks"""
+
+    block = Bottleneck

ml4gw/transforms/__init__.py

@@ -2,5 +2,6 @@ from .pearson import ShiftedPearsonCorrelation
 from .scaler import ChannelWiseScaler
 from .snr_rescaler import SnrRescaler
 from .spectral import SpectralDensity
+from .spectrogram import MultiResolutionSpectrogram
 from .waveforms import WaveformProjector, WaveformSampler
 from .whitening import FixedWhiten, Whiten

ml4gw/transforms/spectrogram.py

@@ -0,0 +1,162 @@
+import warnings
+from typing import Dict, List
+
+import torch
+import torch.nn.functional as F
+from torchaudio.transforms import Spectrogram
+
+
+class MultiResolutionSpectrogram(torch.nn.Module):
+    """
+    Create a batch of multi-resolution spectrograms
+    from a batch of timeseries. Input is expected to
+    have the shape `(B, C, T)`, where `B` is the number
+    of batches, `C` is the number of channels, and `T`
+    is the number of time samples.
+
+    For each timeseries, calculate multiple normalized
+    spectrograms based on the `Spectrogram` `kwargs` given.
+    Combine the spectrograms by taking the maximum value
+    from the nearest time-frequncy bin.
+
+    If the largest number of time bins among the spectrograms
+    is `N` and the largest number of frequency bins is `M`,
+    the output will have dimensions `(B, C, M, N)`
+
+    Args:
+        kernel_length:
+            The length in seconds of the time dimension
+            of the tensor that will be turned into a
+            spectrogram
+        sample_rate:
+            The sample rate of the timeseries in Hz
+        kwargs:
+            Arguments passed in kwargs will used to create
+            `torchaudio.transforms.Spectrogram`s. Each
+            argument should be a list of values. Any list
+            of length greater than 1 should be the same
+            length
+    """
+
+    def __init__(
+        self, kernel_length: float, sample_rate: float, **kwargs
+    ) -> None:
+        super().__init__()
+        self.kernel_size = kernel_length * sample_rate
+        # This method of combination makes sense only when
+        # the spectrograms are normalized, so enforce this
+        if "normalized" in kwargs.keys():
+            if not all(kwargs["normalized"]):
+                raise ValueError(
+                    "Received a value of False for 'normalized'. "
+                    "This method of combination is sensible only for "
+                    "normalized spectrograms."
+                )
+        else:
+            kwargs["normalized"] = [True]
+        self.kwargs = self._check_and_format_kwargs(kwargs)
+
+        self.transforms = torch.nn.ModuleList(
+            [Spectrogram(**k) for k in self.kwargs]
+        )
+
+        dummy_input = torch.ones(int(kernel_length * sample_rate))
+        self.shapes = torch.tensor(
+            [t(dummy_input).shape for t in self.transforms]
+        )
+
+        self.num_freqs = max([shape[0] for shape in self.shapes])
+        self.num_times = max([shape[1] for shape in self.shapes])
+
+        left_pad = torch.zeros(len(self.transforms), dtype=torch.int)
+        top_pad = torch.zeros(len(self.transforms), dtype=torch.int)
+        bottom_pad = torch.tensor(
+            [int(self.num_freqs - shape[0]) for shape in self.shapes]
+        )
+        right_pad = torch.tensor(
+            [int(self.num_times - shape[1]) for shape in self.shapes]
+        )
+        self.register_buffer("left_pad", left_pad)
+        self.register_buffer("top_pad", top_pad)
+        self.register_buffer("bottom_pad", bottom_pad)
+        self.register_buffer("right_pad", right_pad)
+
+        freq_idxs = torch.tensor(
+            [
+                [int(i * shape[0] / self.num_freqs) for shape in self.shapes]
+                for i in range(self.num_freqs)
+            ]
+        )
+        freq_idxs = freq_idxs.repeat(self.num_times, 1, 1).transpose(0, 1)
+        time_idxs = torch.tensor(
+            [
+                [int(i * shape[1] / self.num_times) for shape in self.shapes]
+                for i in range(self.num_times)
+            ]
+        )
+        time_idxs = time_idxs.repeat(self.num_freqs, 1, 1)
+
+        self.register_buffer("freq_idxs", freq_idxs)
+        self.register_buffer("time_idxs", time_idxs)
+
+    def _check_and_format_kwargs(self, kwargs: Dict[str, List]) -> List:
+        lengths = sorted(set([len(v) for v in kwargs.values()]))
+
+        if lengths[-1] > 3:
+            warnings.warn(
+                "Combining too many spectrograms can impede computation time. "
+                "If performance is slower than desired, try reducing the "
+                "number of spectrograms",
+                RuntimeWarning,
+            )
+
+        if len(lengths) > 2 or (len(lengths) == 2 and lengths[0] != 1):
+            raise ValueError(
+                "Spectrogram keyword args should all have the same "
+                f"length or be of length one. Got lengths {lengths}"
+            )
+
+        if len(lengths) == 2:
+            size = lengths[1]
+            kwargs = {k: v * int(size / len(v)) for k, v in kwargs.items()}
+
+        return [dict(zip(kwargs, col)) for col in zip(*kwargs.values())]
+
+    def forward(self, X: torch.Tensor) -> torch.Tensor:
+        """
+        Calculate spectrograms of the input tensor and
+        combine them into a single spectrogram
+
+        Args:
+            X:
+                Batch of multichannel timeseries which will
+                be used to calculate the multi-resolution
+                spectrogram. Should have the shape
+                `(B, C, T)`, where `B` is the number of
+                batches, `C` is the  number of channels,
+                and `T` is the number of time samples.
+        """
+        if X.shape[-1] != self.kernel_size:
+            raise ValueError(
+                "Expected time dimension to be "
+                f"{self.kernel_size} samples long, got input with "
+                f"{X.shape[-1]} samples"
+            )
+
+        spectrograms = [t(X) for t in self.transforms]
+
+        padded_specs = []
+        for spec, left, right, top, bottom in zip(
+            spectrograms,
+            self.left_pad,
+            self.right_pad,
+            self.top_pad,
+            self.bottom_pad,
+        ):
+            padded_specs.append(F.pad(spec, (left, right, top, bottom)))
+
+        padded_specs = torch.stack(padded_specs)
+        remapped_specs = padded_specs[..., self.freq_idxs, self.time_idxs]
+        remapped_specs = torch.diagonal(remapped_specs, dim1=0, dim2=-1)
+
+        return torch.max(remapped_specs, axis=-1)[0]

ml4gw/transforms/whitening.py

@@ -123,7 +123,7 @@ class FixedWhiten(FittableSpectralTransform):
         num_channels: float,
         kernel_length: float,
         sample_rate: float,
-        dtype: torch.dtype = torch.float32,
+        dtype: torch.dtype = torch.float64,
     ) -> None:
         super().__init__()
         self.num_channels = num_channels

ml4gw/waveforms/phenom_d.py

@@ -477,6 +477,7 @@ def rho3_fun(eta, eta2, xi):
 
 def FinalSpin0815(eta, eta2, chi1, chi2):
     Seta = torch.sqrt(1.0 - 4.0 * eta)
+    Seta = torch.nan_to_num(Seta)  # avoid nan around eta = 0.25
     m1 = 0.5 * (1.0 + Seta)
     m2 = 0.5 * (1.0 - Seta)
     m1s = m1 * m1

{ml4gw-0.3.0.dist-info → ml4gw-0.4.0.dist-info}/METADATA

@@ -1,17 +1,19 @@
 Metadata-Version: 2.1
 Name: ml4gw
-Version: 0.3.0
+Version: 0.4.0
 Summary: Tools for training torch models on gravitational wave data
 Author: Alec Gunny
 Author-email: alec.gunny@ligo.org
-Requires-Python: >=3.8,<4.0
+Requires-Python: >=3.8,<3.12
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.8
 Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
-Classifier: Programming Language :: Python :: 3.12
-Requires-Dist: torch (>=1.10,<2.0)
+Requires-Dist: torch (>=1.10,<2.0) ; python_version >= "3.8" and python_version < "3.11"
+Requires-Dist: torch (>=2.0,<3.0) ; python_version >= "3.11"
+Requires-Dist: torchaudio (>=0.13,<0.14) ; python_version >= "3.8" and python_version < "3.11"
+Requires-Dist: torchaudio (>=2.0,<3.0) ; python_version >= "3.11"
 Requires-Dist: torchtyping (>=0.1,<0.2)
 Description-Content-Type: text/markdown
 
@@ -38,8 +40,8 @@ pip install ml4gw torch==1.12.0 --extra-index-url=https://download.pytorch.org/w
 
 ```toml
 [tool.poetry.dependencies]
-python = "^3.8"  # python versions 3.8-3.10 are supported
-ml4gw = "^0.1.0"
+python = "^3.8"  # python versions 3.8-3.11 are supported
+ml4gw = "^0.3.0"
 ```
 
 To build against a specific PyTorch/CUDA combination, consult the PyTorch installation documentation above and specify the `extra-index-url` via the `tool.poetry.source` table in your `pyproject.toml`. For example, to build against CUDA 11.6, you would do something like:
@@ -47,7 +49,7 @@ To build against a specific PyTorch/CUDA combination, consult the PyTorch instal
 ```toml
 [tool.poetry.dependencies]
 python = "^3.8"
-ml4gw = "^0.1.0"
+ml4gw = "^0.3.0"
 torch = {version = "^1.12", source = "torch"}
 
 [[tool.poetry.source]]
@@ -57,6 +59,8 @@ secondary = true
 default = false
 ```
 
+Note: if you are building against CUDA 11.6 or 11.7, make sure that you are using python 3.8, 3.9, or 3.10. Python 3.11 is incompatible with `torchaudio` 0.13, and the following `torchaudio` version is incompatible with CUDA 11.7 and earlier.
+
 ## Use cases
 This library provided utilities for both data iteration and transformation via dataloaders defined in `ml4gw/dataloading` and transform layers exposed in `ml4gw/transforms`. Lower level functions and utilies are defined at the top level of the library and in the `utils` library.
 
@@ -146,3 +150,6 @@ We also strongly encourage ML users in the GW physics space to try their hand at
 For more information about how to get involved, feel free to reach out to [ml4gw@ligo.mit.edu](mailto:ml4gw@ligo.mit.edu) .
 By bringing in new users with new use cases, we hope to develop this library into a truly general-purpose tool which makes DL more accessible for gravitational wave physicists everywhere.
 
+## Funding
+We are grateful for the support of the U.S. National Science Foundation (NSF) Harnessing the Data Revolution (HDR) Institute for <a href="https://a3d3.ai">Accelerating AI Algorithms for Data Driven Discovery (A3D3)</a> under Cooperative Agreement No. <a href="https://www.nsf.gov/awardsearch/showAward?AWD_ID=2117997">PHY-2117997</a>.
+

{ml4gw-0.3.0.dist-info → ml4gw-0.4.0.dist-info}/RECORD

@@ -12,27 +12,32 @@ ml4gw/nn/autoencoder/base.py,sha256=PLr26Cn5DHmgDYX1qj4idfrLehHVeiJqer065ea8_QM,
 ml4gw/nn/autoencoder/convolutional.py,sha256=JTMpTJVdFju9HPPAh9UDdXG1MsFbADrqUIKM8_xg74E,5316
 ml4gw/nn/autoencoder/skip_connection.py,sha256=bOKBLzMqZDh9w8s9G5U93LCESjTSFUHzQGo0hLDOeSk,1304
 ml4gw/nn/autoencoder/utils.py,sha256=whTnWPvdKuVDlxg52azJeM1d9YjiYFWoqIOzJVDGups,326
+ml4gw/nn/norm.py,sha256=9IHZTCCp4zgP7EaGpw1FpAm7o0EU5zu-LYFHKfuLzzw,3250
+ml4gw/nn/resnet/__init__.py,sha256=vBI0IftVP_EYAeDlqomtkGqUYE-RE_S4WNioUhniw9s,64
+ml4gw/nn/resnet/resnet_1d.py,sha256=IQ-EIIzAXd-NWuLwt7JTXLWg5bO3FGJpuFAZwZ78jaI,13218
+ml4gw/nn/resnet/resnet_2d.py,sha256=aK4I0FOZk62JxnYFz0t1O0s5s7J7yRNYSM1flRypvVc,13301
 ml4gw/nn/streaming/__init__.py,sha256=zgjGR2L8t0txXLnil9ceZT0tM8Y2FC8yPxqIKYH0o1A,80
 ml4gw/nn/streaming/online_average.py,sha256=T-wWw7eEufbUVPRNnLAXIq0cedAyJWEE9tdZ6CTi3cs,4561
 ml4gw/nn/streaming/snapshotter.py,sha256=-l_YsWby7ZnEzGIAlLAV2mtR0daLMtLCxovtt4OI3Z0,4432
 ml4gw/spectral.py,sha256=5GfKAV_1vw5yyzTD2u_myjT5jIlAyAHDX6TXj9ynL_o,19021
-ml4gw/transforms/__init__.py,sha256=2V-hZRG-3w3-UczXEvqrBHGpdDbwS2WMBIo6C_7H6Pc,262
+ml4gw/transforms/__init__.py,sha256=t6ZJcq23apqDKhLGM-U5l_bqxJcXFj3riY6cTGY47Gc,314
 ml4gw/transforms/pearson.py,sha256=bJ77lO4wBY6y1R1aESN_bcUEMbc55hWCIaCBdbIj4CY,3133
 ml4gw/transforms/scaler.py,sha256=5VGov0M80NZostRzccViC3HNftx4ZVu0kOKTDmiLrR4,2327
 ml4gw/transforms/snr_rescaler.py,sha256=ocYr6UjpHW7t5TvruV7fyY8KuuDfGOJyvxEulmiFA6o,2275
 ml4gw/transforms/spectral.py,sha256=Vba9199z_ZaxsHWxdpgHB3U216rmGoSyehtvM3R9Z7A,3771
+ml4gw/transforms/spectrogram.py,sha256=R3O8eUB6NHdBFx89v8e_WdJIvXl4qwVeGWZnPyLhHHQ,6024
 ml4gw/transforms/transform.py,sha256=jEr9OFj4u7Wjeh_rpRq90jMpK_TfzcIelbBmt30DxQU,2408
 ml4gw/transforms/waveforms.py,sha256=iyEDSRqK_1zZrxxJenJFbwGUWqbE-alVTXhvjaGl1ww,3060
-ml4gw/transforms/whitening.py,sha256=XDGswhQlmt5IgBvdRypgEhei9SrG1gdYEb6mGHKtO_A,9428
+ml4gw/transforms/whitening.py,sha256=TmvFCCeTOcSEWo5Pt_JQRJ23X5byiJ91q5jHgBRy0rc,9428
 ml4gw/types.py,sha256=XbxunX8zRF95Fp1mZ9jEbixb63bwDQMoayRMMxT9Lzo,429
 ml4gw/utils/interferometer.py,sha256=w_0WkboCJZMKAg-4lhiNGOOkNogAghpT96I0TE5aJ1g,1519
 ml4gw/utils/slicing.py,sha256=Cbwcpk_0hsfN4zczFVM2YbDRjeirA7jFvApM4Jy0U8s,13535
 ml4gw/waveforms/__init__.py,sha256=zjqOKNY4z1A5iPhWTxyhnkLh2robB-obPTtaK-pDUoU,104
 ml4gw/waveforms/generator.py,sha256=4Z6vUEuI84t__3t0DDnXlOyB8R96ynf8xFvtwCGu9JA,1057
-ml4gw/waveforms/phenom_d.py,sha256=pbPUl61e7Z-XthQUKFHmdXRMuhJKi33N1MtGpCJMSN8,38354
+ml4gw/waveforms/phenom_d.py,sha256=pxHk7paW5709Ak29m_DYeQ8kiMLC8wrUnM13flUU36o,38419
 ml4gw/waveforms/phenom_d_data.py,sha256=WA1FBxUp9fo1IQaV_OLJ_5g5gI166mY1FtG9n25he9U,53447
 ml4gw/waveforms/sine_gaussian.py,sha256=WZ6KiVEFSjB9Tv5otJbvI_Yr3341th1Noec_LB9kPOE,3577
 ml4gw/waveforms/taylorf2.py,sha256=x3drvKUMarWI9xHUzMRQhVp1Hh7X-j5WC2bdsbEiVfk,8482
-ml4gw-0.3.0.dist-info/METADATA,sha256=UN66Ak875EILww5qlinoEy5GVc6x7wiRLNKDElTBLrY,5127
-ml4gw-0.3.0.dist-info/WHEEL,sha256=FMvqSimYX_P7y0a7UY-_Mc83r5zkBZsCYPm7Lr0Bsq4,88
-ml4gw-0.3.0.dist-info/RECORD,,
+ml4gw-0.4.0.dist-info/METADATA,sha256=SyBcghsG_wj1TSYjmCnklojeCAH0OQ3uv-DXOoiiuug,5944
+ml4gw-0.4.0.dist-info/WHEEL,sha256=FMvqSimYX_P7y0a7UY-_Mc83r5zkBZsCYPm7Lr0Bsq4,88
+ml4gw-0.4.0.dist-info/RECORD,,