diff --git a/backbones.py b/backbones.py index 9ec6be1..0c42409 100644 --- a/backbones.py +++ b/backbones.py @@ -151,82 +151,27 @@ class AgeSinusoidalEncoding(nn.Module): return output -class GRN2d(nn.Module): - """Global Response Normalization from ConvNeXt V2 for NCHW tensors. +class LiteTimesBackbone2d(nn.Module): + """Cheap local feature extractor for a TimesNet period image.""" - Reference: - Woo et al., "ConvNeXt V2: Co-designing and Scaling ConvNets with - Masked Autoencoders", CVPR 2023. https://arxiv.org/abs/2301.00808 - """ - - def __init__(self, dim: int, eps: float = 1e-6): + def __init__(self, dim: int, kernel_size: int = 5, + expansion: float = 2.0, dropout: float = 0.0): super().__init__() - self.gamma = nn.Parameter(torch.zeros(1, dim, 1, 1)) - self.beta = nn.Parameter(torch.zeros(1, dim, 1, 1)) - self.eps = eps - - def forward(self, x: torch.Tensor) -> torch.Tensor: - gx = torch.norm(x, p=2, dim=(2, 3), keepdim=True) - nx = gx / (gx.mean(dim=1, keepdim=True) + self.eps) - return x + self.gamma * (x * nx) + self.beta - - -class LayerNorm2d(nn.Module): - """Channel-wise LayerNorm for NCHW tensors.""" - - def __init__(self, dim: int, eps: float = 1e-6): - super().__init__() - self.weight = nn.Parameter(torch.ones(dim)) - self.bias = nn.Parameter(torch.zeros(dim)) - self.eps = eps - - def forward(self, x: torch.Tensor) -> torch.Tensor: - x = x.permute(0, 2, 3, 1) - x = F.layer_norm(x, (self.weight.numel(),), self.weight, self.bias, self.eps) - return x.permute(0, 3, 1, 2) - - -class ConvNeXtV2Block2d(nn.Module): - """Lightweight ConvNeXt V2-style 2D block for TimesNet period images. - - This is intentionally a block, not the full ConvNeXt V2 image backbone: - TimesNet's 2D tensors are reshaped time-series period maps rather than - natural images, so aggressive visual downsampling would destroy axis - semantics. - """ - - def __init__( - self, - dim: int, - kernel_size: int = 7, - mlp_ratio: float = 4.0, - dropout: float = 0.0, - ): - super().__init__() - padding = kernel_size // 2 - hidden_dim = int(dim * mlp_ratio) + if kernel_size <= 0 or kernel_size % 2 == 0: + raise ValueError("kernel_size must be a positive odd integer") + if expansion <= 0: + raise ValueError("expansion must be > 0") + hidden_dim = max(dim, int(round(dim * expansion))) self.dwconv = nn.Conv2d( - dim, - dim, - kernel_size=kernel_size, - padding=padding, - groups=dim, + dim, dim, kernel_size=kernel_size, + padding=kernel_size // 2, groups=dim, ) - self.norm = LayerNorm2d(dim) + self.norm = nn.GroupNorm(1, dim) self.pwconv1 = nn.Conv2d(dim, hidden_dim, kernel_size=1) self.act = nn.GELU() - self.grn = GRN2d(hidden_dim) self.pwconv2 = nn.Conv2d(hidden_dim, dim, kernel_size=1) self.drop = nn.Dropout(dropout) - self.reset_parameters() - - def reset_parameters(self) -> None: - nn.init.normal_(self.dwconv.weight, mean=0.0, std=0.02) - nn.init.zeros_(self.dwconv.bias) - nn.init.normal_(self.pwconv1.weight, mean=0.0, std=0.02) - nn.init.zeros_(self.pwconv1.bias) - nn.init.normal_(self.pwconv2.weight, mean=0.0, std=0.02) - nn.init.zeros_(self.pwconv2.bias) + self.layer_scale = nn.Parameter(torch.full((1, dim, 1, 1), 1e-2)) def forward(self, x: torch.Tensor) -> torch.Tensor: residual = x @@ -234,14 +179,12 @@ class ConvNeXtV2Block2d(nn.Module): x = self.norm(x) x = self.pwconv1(x) x = self.act(x) - x = self.grn(x) x = self.pwconv2(x) - x = self.drop(x) - return residual + x + return residual + self.layer_scale * self.drop(x) class TimesNetBlock(nn.Module): - """TimesNet block with ConvNeXt V2-style 2D extraction. + """TimesNet block with lightweight depthwise-separable 2D extraction. The block follows TimesNet's idea: discover dominant periods with FFT, reshape a 1D sequence into period-wise 2D maps, run a 2D convolutional @@ -255,25 +198,27 @@ class TimesNetBlock(nn.Module): def __init__( self, d_model: int, - top_k: int = 3, - n_convnext_blocks: int = 2, - conv_kernel_size: int = 7, - mlp_ratio: float = 4.0, + top_k: int = 2, + n_backbone_blocks: int = 1, + backbone_kernel_size: int = 5, + backbone_expansion: float = 2.0, dropout: float = 0.0, ): super().__init__() if top_k <= 0: raise ValueError(f"top_k must be > 0, got {top_k}") + if n_backbone_blocks <= 0: + raise ValueError("n_backbone_blocks must be > 0") self.top_k = top_k self.norm = nn.LayerNorm(d_model) self.extractor = nn.Sequential(*[ - ConvNeXtV2Block2d( + LiteTimesBackbone2d( dim=d_model, - kernel_size=conv_kernel_size, - mlp_ratio=mlp_ratio, + kernel_size=backbone_kernel_size, + expansion=backbone_expansion, dropout=dropout, ) - for _ in range(n_convnext_blocks) + for _ in range(n_backbone_blocks) ]) def _select_periods(self, x: torch.Tensor) -> tuple[list[int], torch.Tensor]: @@ -326,10 +271,10 @@ class TimesNetEncoder(nn.Module): input_dim: int, d_model: int, n_layers: int = 2, - top_k: int = 3, - n_convnext_blocks: int = 2, - conv_kernel_size: int = 7, - mlp_ratio: float = 4.0, + top_k: int = 2, + n_backbone_blocks: int = 1, + backbone_kernel_size: int = 5, + backbone_expansion: float = 2.0, dropout: float = 0.0, append_observed_mask: bool = True, ): @@ -342,9 +287,9 @@ class TimesNetEncoder(nn.Module): TimesNetBlock( d_model=d_model, top_k=top_k, - n_convnext_blocks=n_convnext_blocks, - conv_kernel_size=conv_kernel_size, - mlp_ratio=mlp_ratio, + n_backbone_blocks=n_backbone_blocks, + backbone_kernel_size=backbone_kernel_size, + backbone_expansion=backbone_expansion, dropout=dropout, ) for _ in range(n_layers) @@ -415,10 +360,10 @@ class TimesNetExposureEncoder(nn.Module): monthly_input_dim: int = 2, d_model: int | None = None, n_layers: int = 2, - top_k: int = 3, - n_convnext_blocks: int = 2, - conv_kernel_size: int = 7, - mlp_ratio: float = 4.0, + top_k: int = 2, + n_backbone_blocks: int = 1, + backbone_kernel_size: int = 5, + backbone_expansion: float = 2.0, dropout: float = 0.0, use_gate: bool = True, ): @@ -429,9 +374,9 @@ class TimesNetExposureEncoder(nn.Module): d_model=d_model, n_layers=n_layers, top_k=top_k, - n_convnext_blocks=n_convnext_blocks, - conv_kernel_size=conv_kernel_size, - mlp_ratio=mlp_ratio, + n_backbone_blocks=n_backbone_blocks, + backbone_kernel_size=backbone_kernel_size, + backbone_expansion=backbone_expansion, dropout=dropout, append_observed_mask=True, ) @@ -440,9 +385,9 @@ class TimesNetExposureEncoder(nn.Module): d_model=d_model, n_layers=n_layers, top_k=top_k, - n_convnext_blocks=n_convnext_blocks, - conv_kernel_size=conv_kernel_size, - mlp_ratio=mlp_ratio, + n_backbone_blocks=n_backbone_blocks, + backbone_kernel_size=backbone_kernel_size, + backbone_expansion=backbone_expansion, dropout=dropout, append_observed_mask=True, ) @@ -485,10 +430,10 @@ class TimesNetSequenceDecoder(nn.Module): latent_dim: int, d_model: int, n_layers: int = 2, - top_k: int = 3, - n_convnext_blocks: int = 2, - conv_kernel_size: int = 7, - mlp_ratio: float = 4.0, + top_k: int = 2, + n_backbone_blocks: int = 1, + backbone_kernel_size: int = 5, + backbone_expansion: float = 2.0, dropout: float = 0.0, ): super().__init__() @@ -498,9 +443,9 @@ class TimesNetSequenceDecoder(nn.Module): TimesNetBlock( d_model=d_model, top_k=top_k, - n_convnext_blocks=n_convnext_blocks, - conv_kernel_size=conv_kernel_size, - mlp_ratio=mlp_ratio, + n_backbone_blocks=n_backbone_blocks, + backbone_kernel_size=backbone_kernel_size, + backbone_expansion=backbone_expansion, dropout=dropout, ) for _ in range(n_layers) @@ -538,10 +483,10 @@ class TimesNetExposureAutoencoder(nn.Module): monthly_input_dim: int = 2, d_model: int | None = None, n_layers: int = 2, - top_k: int = 3, - n_convnext_blocks: int = 2, - conv_kernel_size: int = 7, - mlp_ratio: float = 4.0, + top_k: int = 2, + n_backbone_blocks: int = 1, + backbone_kernel_size: int = 5, + backbone_expansion: float = 2.0, dropout: float = 0.0, ): super().__init__() @@ -553,9 +498,9 @@ class TimesNetExposureAutoencoder(nn.Module): d_model=d_model, n_layers=n_layers, top_k=top_k, - n_convnext_blocks=n_convnext_blocks, - conv_kernel_size=conv_kernel_size, - mlp_ratio=mlp_ratio, + n_backbone_blocks=n_backbone_blocks, + backbone_kernel_size=backbone_kernel_size, + backbone_expansion=backbone_expansion, dropout=dropout, use_gate=True, ) @@ -564,9 +509,9 @@ class TimesNetExposureAutoencoder(nn.Module): d_model=d_model, n_layers=n_layers, top_k=top_k, - n_convnext_blocks=n_convnext_blocks, - conv_kernel_size=conv_kernel_size, - mlp_ratio=mlp_ratio, + n_backbone_blocks=n_backbone_blocks, + backbone_kernel_size=backbone_kernel_size, + backbone_expansion=backbone_expansion, dropout=dropout, ) self.encoder = TimesNetExposureEncoder(**encoder_kwargs) diff --git a/evaluate_auc.py b/evaluate_auc.py index 93a7d1f..0dcc2c3 100644 --- a/evaluate_auc.py +++ b/evaluate_auc.py @@ -324,12 +324,12 @@ def build_model_from_dataset(args: argparse.Namespace, cfg: Dict[str, Any], data dist_mode=str(cfg_get(args, cfg, "dist_mode", "exponential")), dropout=float(cfg_get(args, cfg, "dropout", 0.0)), use_exposure_encoder=bool(cfg_get(args, cfg, "use_exposure_encoder", False)), - exposure_d_model=cfg_get(args, cfg, "exposure_d_model", None), + exposure_d_model=cfg_get(args, cfg, "exposure_d_model", 64), exposure_n_layers=int(cfg_get(args, cfg, "exposure_n_layers", 2)), - exposure_top_k=int(cfg_get(args, cfg, "exposure_top_k", 3)), - exposure_n_convnext_blocks=int(cfg_get(args, cfg, "exposure_n_convnext_blocks", 2)), - exposure_conv_kernel_size=int(cfg_get(args, cfg, "exposure_conv_kernel_size", 7)), - exposure_mlp_ratio=float(cfg_get(args, cfg, "exposure_mlp_ratio", 4.0)), + exposure_top_k=int(cfg_get(args, cfg, "exposure_top_k", 2)), + exposure_n_backbone_blocks=int(cfg_get(args, cfg, "exposure_n_backbone_blocks", 1)), + exposure_backbone_kernel_size=int(cfg_get(args, cfg, "exposure_backbone_kernel_size", 5)), + exposure_backbone_expansion=float(cfg_get(args, cfg, "exposure_backbone_expansion", 2.0)), exposure_use_gate=bool(cfg_get(args, cfg, "exposure_use_gate", True)), ) diff --git a/models.py b/models.py index c9e1467..72cb53e 100644 --- a/models.py +++ b/models.py @@ -33,12 +33,12 @@ class DeepHealth(nn.Module): use_exposure_encoder: bool = False, exposure_daily_input_dim: int = 4, exposure_monthly_input_dim: int = 2, - exposure_d_model: int | None = None, + exposure_d_model: int | None = 64, exposure_n_layers: int = 2, - exposure_top_k: int = 3, - exposure_n_convnext_blocks: int = 2, - exposure_conv_kernel_size: int = 7, - exposure_mlp_ratio: float = 4.0, + exposure_top_k: int = 2, + exposure_n_backbone_blocks: int = 1, + exposure_backbone_kernel_size: int = 5, + exposure_backbone_expansion: float = 2.0, exposure_use_gate: bool = True, ): super().__init__() @@ -64,9 +64,9 @@ class DeepHealth(nn.Module): d_model=exposure_d_model, n_layers=exposure_n_layers, top_k=exposure_top_k, - n_convnext_blocks=exposure_n_convnext_blocks, - conv_kernel_size=exposure_conv_kernel_size, - mlp_ratio=exposure_mlp_ratio, + n_backbone_blocks=exposure_n_backbone_blocks, + backbone_kernel_size=exposure_backbone_kernel_size, + backbone_expansion=exposure_backbone_expansion, dropout=dropout, use_gate=exposure_use_gate, ) diff --git a/train_exposure_autoencoder.py b/train_exposure_autoencoder.py index c3204fa..afdd203 100644 --- a/train_exposure_autoencoder.py +++ b/train_exposure_autoencoder.py @@ -1,4 +1,4 @@ -"""Pretrain a TimesNet + ConvNeXtV2 autoencoder on training-set exposure.""" +"""Pretrain a lightweight TimesNet autoencoder on training-set exposure.""" from __future__ import annotations import argparse @@ -46,7 +46,7 @@ class ExposureWindowDataset(Dataset): def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( - description="Pretrain a TimesNet + ConvNeXtV2 exposure autoencoder" + description="Pretrain a lightweight TimesNet exposure autoencoder" ) parser.add_argument("--exposure_cache_dir", required=True) parser.add_argument("--train_eid_file", default="ukb_train_eid.csv") @@ -54,12 +54,12 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--runs_root", default="runs") parser.add_argument("--seed", type=int, default=42) parser.add_argument("--n_embd", type=int, default=120) - parser.add_argument("--d_model", type=int, default=None) + parser.add_argument("--d_model", type=int, default=64) parser.add_argument("--n_layers", type=int, default=2) - parser.add_argument("--top_k", type=int, default=3) - parser.add_argument("--n_convnext_blocks", type=int, default=2) - parser.add_argument("--conv_kernel_size", type=int, default=7) - parser.add_argument("--mlp_ratio", type=float, default=4.0) + parser.add_argument("--top_k", type=int, default=2) + parser.add_argument("--n_backbone_blocks", type=int, default=1) + parser.add_argument("--backbone_kernel_size", type=int, default=5) + parser.add_argument("--backbone_expansion", type=float, default=2.0) parser.add_argument("--dropout", type=float, default=0.0) parser.add_argument("--mask_ratio", type=float, default=0.25) parser.add_argument("--batch_size", type=int, default=16) @@ -288,8 +288,9 @@ def main() -> None: ) model = TimesNetExposureAutoencoder( n_embd=args.n_embd, d_model=args.d_model, n_layers=args.n_layers, - top_k=args.top_k, n_convnext_blocks=args.n_convnext_blocks, - conv_kernel_size=args.conv_kernel_size, mlp_ratio=args.mlp_ratio, + top_k=args.top_k, n_backbone_blocks=args.n_backbone_blocks, + backbone_kernel_size=args.backbone_kernel_size, + backbone_expansion=args.backbone_expansion, dropout=args.dropout, ).to(device) model = maybe_wrap_data_parallel(model, args, device, logger) @@ -349,8 +350,8 @@ def main() -> None: "model_config": { key: config[key] for key in ( "n_embd", "d_model", "n_layers", "top_k", - "n_convnext_blocks", "conv_kernel_size", - "mlp_ratio", "dropout", + "n_backbone_blocks", "backbone_kernel_size", + "backbone_expansion", "dropout", ) }, "normalization": { diff --git a/train_next_step.py b/train_next_step.py index 8dbe89b..7d13644 100644 --- a/train_next_step.py +++ b/train_next_step.py @@ -129,12 +129,12 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--dropout", type=float, default=0.0) parser.add_argument("--exposure_cache_dir", type=str, default=None) parser.add_argument("--mask_onset_exposure", action="store_true") - parser.add_argument("--exposure_d_model", type=int, default=None) + parser.add_argument("--exposure_d_model", type=int, default=64) parser.add_argument("--exposure_n_layers", type=int, default=2) - parser.add_argument("--exposure_top_k", type=int, default=3) - parser.add_argument("--exposure_n_convnext_blocks", type=int, default=2) - parser.add_argument("--exposure_conv_kernel_size", type=int, default=7) - parser.add_argument("--exposure_mlp_ratio", type=float, default=4.0) + parser.add_argument("--exposure_top_k", type=int, default=2) + parser.add_argument("--exposure_n_backbone_blocks", type=int, default=1) + parser.add_argument("--exposure_backbone_kernel_size", type=int, default=5) + parser.add_argument("--exposure_backbone_expansion", type=float, default=2.0) parser.add_argument("--no_exposure_gate", action="store_true") parser.add_argument("--target_mode", type=str, default="uts", choices=["delphi2m", "uts"]) @@ -280,9 +280,9 @@ def build_model(args: argparse.Namespace, dataset: HealthDataset) -> DeepHealth: exposure_d_model=args.exposure_d_model, exposure_n_layers=args.exposure_n_layers, exposure_top_k=args.exposure_top_k, - exposure_n_convnext_blocks=args.exposure_n_convnext_blocks, - exposure_conv_kernel_size=args.exposure_conv_kernel_size, - exposure_mlp_ratio=args.exposure_mlp_ratio, + exposure_n_backbone_blocks=args.exposure_n_backbone_blocks, + exposure_backbone_kernel_size=args.exposure_backbone_kernel_size, + exposure_backbone_expansion=args.exposure_backbone_expansion, exposure_use_gate=not args.no_exposure_gate, ) @@ -490,9 +490,9 @@ def build_metadata( "exposure_d_model": args.exposure_d_model, "exposure_n_layers": int(args.exposure_n_layers), "exposure_top_k": int(args.exposure_top_k), - "exposure_n_convnext_blocks": int(args.exposure_n_convnext_blocks), - "exposure_conv_kernel_size": int(args.exposure_conv_kernel_size), - "exposure_mlp_ratio": float(args.exposure_mlp_ratio), + "exposure_n_backbone_blocks": int(args.exposure_n_backbone_blocks), + "exposure_backbone_kernel_size": int(args.exposure_backbone_kernel_size), + "exposure_backbone_expansion": float(args.exposure_backbone_expansion), "exposure_use_gate": not bool(args.no_exposure_gate), "num_workers": int(args.num_workers), "prefetch_factor": int(args.prefetch_factor),