ml4gw 0.3.0__tar.gz → 0.4.0__tar.gz

{ml4gw-0.3.0 → ml4gw-0.4.0}/PKG-INFO

@@ -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 → ml4gw-0.4.0}/README.md

@@ -21,8 +21,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:
@@ -30,7 +30,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]]
@@ -40,6 +40,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.
 
@@ -128,3 +130,6 @@ We encourage users who encounter these difficulties to file issues on GitHub, an
 We also strongly encourage ML users in the GW physics space to try their hand at working on these issues and joining on as collaborators!
 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.4.0/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-0.4.0/ml4gw/nn/resnet/__init__.py

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

ml4gw-0.4.0/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