| """ |
| Data augmentations for SSL pretraining (SimCLR, DINO). |
| """ |
| import torch |
| import torch.nn.functional as F |
| from typing import Tuple |
|
|
|
|
| @torch.no_grad() |
| def random_time_crop( |
| x: torch.Tensor, |
| ratio: Tuple[float, float] | float = (0.6, 0.9), |
| *, |
| resize_back: bool = True, |
| align_to: int | None = 40 |
| ) -> torch.Tensor: |
| """ |
| Randomly crop a contiguous sub-sequence per sample, optionally resize back to original T. |
| |
| Args: |
| x: (B, C, T) |
| ratio: crop length ratio in [low, high] or a float |
| resize_back: if True, linearly interpolate the cropped view back to length T |
| align_to: if not None, crop length is rounded to a multiple of align_to (>= align_to) |
| """ |
| assert x.dim() == 3, f"expected (B,C,T), got {tuple(x.shape)}" |
| B, C, T = x.shape |
| dev = x.device |
|
|
| def _sample_L() -> int: |
| if isinstance(ratio, (tuple, list)): |
| a, b = float(ratio[0]), float(ratio[1]) |
| r = torch.empty((), device=dev).uniform_(a, b).item() |
| else: |
| r = float(ratio) |
| L = max(2, int(round(T * r))) |
| if align_to and align_to > 1: |
| L = max(align_to, int(round(L / align_to)) * align_to) |
| return min(L, T) |
|
|
| Ls = [_sample_L() for _ in range(B)] |
| outs = [] |
| for b in range(B): |
| L = Ls[b] |
| max_start = max(0, T - L) |
| s = int(torch.randint(0, max_start + 1, (1,), device=dev).item()) |
| v = x[b, :, s:s+L] |
| if resize_back and v.shape[-1] != T: |
| v = F.interpolate(v[None], size=T, mode="linear", align_corners=False)[0] |
| outs.append(v) |
| return torch.stack(outs, dim=0) |
|
|
|
|
| @torch.no_grad() |
| def channel_dropout( |
| x: torch.Tensor, |
| drop_prob: float = 0.2, |
| min_keep: int = 1 |
| ) -> torch.Tensor: |
| """ |
| Drop entire channels to zero with probability drop_prob (per sample, per channel). |
| Ensures at least `min_keep` channels remain active in each sample. |
| |
| Args: |
| x: (B, C, T) |
| drop_prob: probability to drop each channel |
| min_keep: minimum number of channels to keep per sample |
| """ |
| assert x.dim() == 3 |
| B, C, T = x.shape |
| mask = (torch.rand(B, C, 1, device=x.device, dtype=x.dtype) > drop_prob).to(x.dtype) |
| |
| |
| keep = mask.sum(dim=1, keepdim=True) |
| need = (keep < min_keep).squeeze(-1).squeeze(-1) |
| if need.any(): |
| for b in torch.where(need)[0]: |
| idx = torch.randperm(C, device=x.device)[:min_keep] |
| mask[b, idx, 0] = 1.0 |
| |
| return x * mask |
|
|
