rc-foundry 0.1.1__py3-none-any.whl

foundry/training/EMA.py

@@ -0,0 +1,67 @@
+from collections import OrderedDict
+from copy import deepcopy
+
+import torch
+import torch.nn as nn
+
+
+class EMA(nn.Module):
+    # TODO: Rename shadow to `ema_model` to better match convention
+    def __init__(self, model: nn.Module, decay: float):
+        """Initialize the Exponential Moving Average (EMA) module.
+
+        EMA maintains a shadow model that slowly tracks the weight of the original model.
+
+        Args:
+            model: The original model.
+            decay: The decay rate of the EMA. The shadow model will be updated with the formula:
+                shadow_variable -= (1 - decay) * (shadow_variable - variable)
+        """
+        super().__init__()
+        self.decay = decay
+
+        self.model = model
+        self.shadow = deepcopy(self.model)
+
+        # Detach the shadow model from the computation graph
+        for param in self.shadow.parameters():
+            param.detach_()
+
+    @torch.no_grad()
+    def update(self):
+        """Update the shadow model using the weight of the original model and the decay rate."""
+        if not self.training:
+            raise RuntimeError("EMA update should only be called during training")
+
+        # ... get the model and shadow parameters
+        model_params = OrderedDict(self.model.named_parameters())
+        shadow_params = OrderedDict(self.shadow.named_parameters())
+
+        # ... ensure that both models have the same set of keys
+        assert model_params.keys() == shadow_params.keys()
+
+        for name, param in model_params.items():
+            # Update the shadow model with the formula:
+            # shadow_variable -= (1 - decay) * (shadow_variable - variable)
+            # Reference: https://www.tensorflow.org/api_docs/python/tf/train/ExponentialMovingAverage
+            if param.requires_grad:
+                shadow_params[name].sub_(
+                    (1.0 - self.decay) * (shadow_params[name] - param)
+                )
+
+        # ... and do the same with the buffers (e.g,. objects that are part of the module state but not trainable parameters)
+        model_buffers = OrderedDict(self.model.named_buffers())
+        shadow_buffers = OrderedDict(self.shadow.named_buffers())
+
+        assert model_buffers.keys() == shadow_buffers.keys()
+
+        for name, buffer in model_buffers.items():
+            #  ... copy the buffers from the model to the shadow
+            shadow_buffers[name].copy_(buffer)
+
+    def forward(self, *args, **kwargs):
+        """Dynamic dispatch to the correct model (model or shadow)."""
+        if self.training:
+            return self.model(*args, **kwargs)
+        else:
+            return self.shadow(*args, **kwargs)

foundry/training/checkpoint.py

@@ -0,0 +1,61 @@
+"""Utilities for gradient checkpointing.
+
+References:
+  * `PyTorch Checkpoint Documentation`_
+
+  .. _PyTorch Checkpoint Documentation: https://pytorch.org/docs/stable/checkpoint.html
+"""
+
+import torch
+from torch.utils.checkpoint import checkpoint
+
+
+def create_custom_forward(module, **kwargs):
+    """Create a custom forward function for gradient checkpointing with fixed kwargs.
+
+    Enables passing keyword arguments to a module when using PyTorch's checkpoint function,
+    which only accepts positional arguments for the function to be checkpointed.
+
+    Args:
+      module: The callable (typically a nn.Module) to wrap.
+      **kwargs: Keyword arguments to pass to the module during forward.
+
+    Returns:
+      A callable that accepts only positional arguments and forwards them along
+      with the fixed kwargs to the original module.
+    """
+
+    def custom_forward(*inputs):
+        return module(*inputs, **kwargs)
+
+    return custom_forward
+
+
+def activation_checkpointing(function):
+    """Decorator to enable gradient checkpointing for a function during training.
+
+    Args:
+      function: The function to apply gradient checkpointing to.
+
+    Returns:
+      Wrapped function that conditionally applies checkpointing based on gradient state.
+
+    Examples:
+      Apply to a forward pass method::
+
+        @activation_checkpointing
+        def forward(self, x, mask=None):
+            return self.layer(x, mask)
+
+    Notes:
+      Uses ``use_reentrant=False`` for compatibility with recent PyTorch versions.
+    """
+
+    def wrapper(*args, **kwargs):
+        if torch.is_grad_enabled():
+            return checkpoint(
+                create_custom_forward(function, **kwargs), *args, use_reentrant=False
+            )
+        return function(*args, **kwargs)
+
+    return wrapper

foundry/training/schedulers.py

@@ -0,0 +1,91 @@
+from dataclasses import dataclass
+
+from torch.optim.lr_scheduler import LRScheduler, _LRScheduler
+from torch.optim.optimizer import Optimizer
+
+
+class AF3Scheduler(_LRScheduler):
+    """Implements a two-phase learning rate schedule a-la AF-3:
+        1. The base learning rate is 1.8 · 10^−3, which is linearly increased from 0 over the first 1,000 steps.
+        2. The learning rate is then decreased by a factor of 0.95 every 50,000 steps.
+
+    From the AF-3 Supplement, Section 5.4:
+    >  "For training we use the Adam optimizer with parameters β1 = 0.9, β2 = 0.95, ϵ = 10^−8. The base learning rate
+        is 1.8 · 10^−3, which is linearly increased from 0 over the first 1,000 steps. The learning rate is then decreased
+        by a factor of 0.95 every 5 · 10^4 steps."
+
+    References:
+        - AF-3 Supplement
+    """
+
+    def __init__(
+        self,
+        optimizer: Optimizer,
+        base_lr: float = 1.8e-3,
+        warmup_steps: int = 1000,
+        decay_factor: float = 0.95,
+        decay_steps: int = 50000,
+        last_epoch: int = -1,
+    ) -> None:
+        """Initializes a new instance of AF3LRScheduler.
+
+        Note that the "last_epoch" value is incremented every time we call `scheduler.step()`
+        method; we name it "epoch" to follow the PyTorch convention.
+
+        Args:
+            optimizer (Optimizer): Wrapped optimizer.
+            base_lr (float): The base learning rate after warmup (which will then be decayed).
+            warmup_steps (int): Number of steps for linear warmup.
+            decay_factor (float): Factor by which the learning rate is multiplied every decay_steps.
+            decay_steps (int): Number of steps between each decay.
+            last_epoch (int): The index of the last epoch. Default: -1.
+        """
+        self.base_lr = base_lr
+        self.warmup_steps = warmup_steps
+        self.decay_factor = decay_factor
+        self.decay_steps = decay_steps
+        super(AF3Scheduler, self).__init__(optimizer, last_epoch)
+
+    def get_lr(self) -> list[float]:
+        if self.last_epoch < self.warmup_steps:
+            # Linear warmup
+            return [
+                self.base_lr * (self.last_epoch / self.warmup_steps)
+                for _ in self.optimizer.param_groups
+            ]
+        else:
+            # Decay after warmup
+            num_decays = (self.last_epoch - self.warmup_steps) // self.decay_steps
+            return [
+                self.base_lr * (self.decay_factor**num_decays)
+                for _ in self.optimizer.param_groups
+            ]
+
+
+@dataclass
+class SchedulerConfig:
+    """Flexible configuration for a learning rate scheduler.
+
+    Modeled on the PyTorch Lightning scheduler configuration.
+
+    Attributes:
+        scheduler (LRScheduler): The learning rate scheduler instance. Must inherit from `torch.optim.lr_scheduler.LRScheduler`.
+        interval (str): The interval at which to apply the scheduler, typically "epoch" or "step". Defaults to "step".
+        frequency (int): The frequency of applying the scheduler. For example, a frequency of 1 means the scheduler is applied every epoch. Defaults to 1.
+    """
+
+    scheduler: LRScheduler = None
+    interval: str = "step"
+    frequency: int = 1
+
+    def state_dict(self) -> dict:
+        return {
+            "scheduler": self.scheduler.state_dict(),
+            "interval": self.interval,
+            "frequency": self.frequency,
+        }
+
+    def load_state_dict(self, state_dict: dict) -> None:
+        self.scheduler.load_state_dict(state_dict["scheduler"])
+        self.interval = state_dict["interval"]
+        self.frequency = state_dict["frequency"]

foundry/utils/alignment.py

@@ -0,0 +1,86 @@
+import logging
+
+import torch
+
+logger = logging.getLogger(__name__)
+
+
+def weighted_rigid_align(
+    X_L,  # [B, L, 3]
+    X_gt_L,  # [B, L, 3]
+    X_exists_L=None,  # [L]
+    w_L=None,  # [B, L]
+):
+    """
+    Weighted rigid body alignment of X_gt_L onto X_L with weights w_L
+    Allows for "moving target" ground truth that is se3 invariant
+    Following algorithm 28 in AF3 paper
+    Returns:
+      X_align_L: [B, L, 3]
+    """
+    assert X_L.shape == X_gt_L.shape
+    assert X_L.shape[:-1] == w_L.shape
+
+    if X_exists_L is None:
+        X_exists_L = torch.ones((X_L.shape[-2]), dtype=torch.bool)
+    if w_L is None:
+        w_L = torch.ones_like(X_L[..., 0])
+    else:
+        w_L = w_L.to(torch.float32)
+
+    # Assert `X_exists_L` is a boolean mask
+    assert (
+        X_exists_L.dtype == torch.bool
+    ), "X_exists_L should be a boolean mask! Otherwise, the alignment will be incorrect (silent failure)!"
+
+    X_resolved = X_L[:, X_exists_L]
+    X_gt_resolved = X_gt_L[:, X_exists_L]
+    w_resolved = w_L[:, X_exists_L]
+    u_X = torch.sum(X_resolved * w_resolved.unsqueeze(-1), dim=-2) / torch.sum(
+        w_resolved, dim=-1, keepdim=True
+    )
+    u_X_gt = torch.sum(X_gt_resolved * w_resolved.unsqueeze(-1), dim=-2) / torch.sum(
+        w_resolved, dim=-1, keepdim=True
+    )
+
+    X_resolved = X_resolved - u_X.unsqueeze(-2)
+    X_gt_resolved = X_gt_resolved - u_X_gt.unsqueeze(-2)
+
+    # Computation of the covariance matrix
+    C = torch.einsum("bji,bjk->bik", w_resolved[..., None] * X_gt_resolved, X_resolved)
+
+    U, S, V = torch.linalg.svd(C)
+
+    R = U @ V
+    B, _, _ = X_L.shape
+    F = torch.eye(3, 3, device=X_L.device)[None].tile(
+        (
+            B,
+            1,
+            1,
+        )
+    )
+
+    F[..., -1, -1] = torch.sign(torch.linalg.det(R))
+    R = U @ F @ V
+
+    X_gt_L = X_gt_L - u_X_gt.unsqueeze(-2)
+    X_align_L = X_gt_L @ R + u_X.unsqueeze(-2)
+
+    return X_align_L.detach()
+
+
+def get_rmsd(xyz1, xyz2, eps=1e-4):
+    L = xyz1.shape[-2]
+    rmsd = torch.sqrt(torch.sum((xyz2 - xyz1) * (xyz2 - xyz1), axis=(-1, -2)) / L + eps)
+    return rmsd
+
+
+def superimpose(xyz1, xyz2, mask, eps=1e-4):
+    """
+    Superimpose xyz1 onto xyz2 using mask
+    """
+    L = xyz1.shape[-2]
+    assert mask.shape == (L,)
+    assert xyz1.shape == xyz2.shape
+    assert mask.dtype == torch.bool