JiaruiLi/DeepHealthExpo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Lighten TimesNet exposure backbone

Browse Source
This commit is contained in:
Jiarui Li
2026-07-09 14:23:28 +08:00
parent 54fedc620b
commit f7fb6b7718
5 changed files with 95 additions and 149 deletions
Download Patch File Download Diff File Expand all files Collapse all files

173
backbones.py
View File

@@ -151,82 +151,27 @@ class AgeSinusoidalEncoding(nn.Module):
return output return output
class GRN2d(nn.Module): class LiteTimesBackbone2d(nn.Module):
"""Global Response Normalization from ConvNeXt V2 for NCHW tensors. """Cheap local feature extractor for a TimesNet period image."""
Reference: def __init__(self, dim: int, kernel_size: int = 5,
Woo et al., "ConvNeXt V2: Co-designing and Scaling ConvNets with expansion: float = 2.0, dropout: float = 0.0):
Masked Autoencoders", CVPR 2023. https://arxiv.org/abs/2301.00808
"""
def __init__(self, dim: int, eps: float = 1e-6):
super().__init__() super().__init__()
self.gamma = nn.Parameter(torch.zeros(1, dim, 1, 1)) if kernel_size <= 0 or kernel_size % 2 == 0:
self.beta = nn.Parameter(torch.zeros(1, dim, 1, 1)) raise ValueError("kernel_size must be a positive odd integer")
self.eps = eps if expansion <= 0:
raise ValueError("expansion must be > 0")
def forward(self, x: torch.Tensor) -> torch.Tensor: hidden_dim = max(dim, int(round(dim * expansion)))
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)
self.dwconv = nn.Conv2d( self.dwconv = nn.Conv2d(
dim, dim, dim, kernel_size=kernel_size,
dim, padding=kernel_size // 2, groups=dim,
kernel_size=kernel_size,
padding=padding,
groups=dim,
) )
self.norm = LayerNorm2d(dim) self.norm = nn.GroupNorm(1, dim)
self.pwconv1 = nn.Conv2d(dim, hidden_dim, kernel_size=1) self.pwconv1 = nn.Conv2d(dim, hidden_dim, kernel_size=1)
self.act = nn.GELU() self.act = nn.GELU()
self.grn = GRN2d(hidden_dim)
self.pwconv2 = nn.Conv2d(hidden_dim, dim, kernel_size=1) self.pwconv2 = nn.Conv2d(hidden_dim, dim, kernel_size=1)
self.drop = nn.Dropout(dropout) self.drop = nn.Dropout(dropout)
self.reset_parameters() self.layer_scale = nn.Parameter(torch.full((1, dim, 1, 1), 1e-2))
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)
def forward(self, x: torch.Tensor) -> torch.Tensor: def forward(self, x: torch.Tensor) -> torch.Tensor:
residual = x residual = x
@@ -234,14 +179,12 @@ class ConvNeXtV2Block2d(nn.Module):
x = self.norm(x) x = self.norm(x)
x = self.pwconv1(x) x = self.pwconv1(x)
x = self.act(x) x = self.act(x)
x = self.grn(x)
x = self.pwconv2(x) x = self.pwconv2(x)
x = self.drop(x) return residual + self.layer_scale * self.drop(x)
return residual + x
class TimesNetBlock(nn.Module): 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, The block follows TimesNet's idea: discover dominant periods with FFT,
reshape a 1D sequence into period-wise 2D maps, run a 2D convolutional reshape a 1D sequence into period-wise 2D maps, run a 2D convolutional
@@ -255,25 +198,27 @@ class TimesNetBlock(nn.Module):
def __init__( def __init__(
self, self,
d_model: int, d_model: int,
top_k: int = 3, top_k: int = 2,
n_convnext_blocks: int = 2, n_backbone_blocks: int = 1,
conv_kernel_size: int = 7, backbone_kernel_size: int = 5,
mlp_ratio: float = 4.0, backbone_expansion: float = 2.0,
dropout: float = 0.0, dropout: float = 0.0,
): ):
super().__init__() super().__init__()
if top_k <= 0: if top_k <= 0:
raise ValueError(f"top_k must be > 0, got {top_k}") 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.top_k = top_k
self.norm = nn.LayerNorm(d_model) self.norm = nn.LayerNorm(d_model)
self.extractor = nn.Sequential(*[ self.extractor = nn.Sequential(*[
ConvNeXtV2Block2d( LiteTimesBackbone2d(
dim=d_model, dim=d_model,
kernel_size=conv_kernel_size, kernel_size=backbone_kernel_size,
mlp_ratio=mlp_ratio, expansion=backbone_expansion,
dropout=dropout, 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]: def _select_periods(self, x: torch.Tensor) -> tuple[list[int], torch.Tensor]:
@@ -326,10 +271,10 @@ class TimesNetEncoder(nn.Module):
input_dim: int, input_dim: int,
d_model: int, d_model: int,
n_layers: int = 2, n_layers: int = 2,
top_k: int = 3, top_k: int = 2,
n_convnext_blocks: int = 2, n_backbone_blocks: int = 1,
conv_kernel_size: int = 7, backbone_kernel_size: int = 5,
mlp_ratio: float = 4.0, backbone_expansion: float = 2.0,
dropout: float = 0.0, dropout: float = 0.0,
append_observed_mask: bool = True, append_observed_mask: bool = True,
): ):
@@ -342,9 +287,9 @@ class TimesNetEncoder(nn.Module):
TimesNetBlock( TimesNetBlock(
d_model=d_model, d_model=d_model,
top_k=top_k, top_k=top_k,
n_convnext_blocks=n_convnext_blocks, n_backbone_blocks=n_backbone_blocks,
conv_kernel_size=conv_kernel_size, backbone_kernel_size=backbone_kernel_size,
mlp_ratio=mlp_ratio, backbone_expansion=backbone_expansion,
dropout=dropout, dropout=dropout,
) )
for _ in range(n_layers) for _ in range(n_layers)
@@ -415,10 +360,10 @@ class TimesNetExposureEncoder(nn.Module):
monthly_input_dim: int = 2, monthly_input_dim: int = 2,
d_model: int | None = None, d_model: int | None = None,
n_layers: int = 2, n_layers: int = 2,
top_k: int = 3, top_k: int = 2,
n_convnext_blocks: int = 2, n_backbone_blocks: int = 1,
conv_kernel_size: int = 7, backbone_kernel_size: int = 5,
mlp_ratio: float = 4.0, backbone_expansion: float = 2.0,
dropout: float = 0.0, dropout: float = 0.0,
use_gate: bool = True, use_gate: bool = True,
): ):
@@ -429,9 +374,9 @@ class TimesNetExposureEncoder(nn.Module):
d_model=d_model, d_model=d_model,
n_layers=n_layers, n_layers=n_layers,
top_k=top_k, top_k=top_k,
n_convnext_blocks=n_convnext_blocks, n_backbone_blocks=n_backbone_blocks,
conv_kernel_size=conv_kernel_size, backbone_kernel_size=backbone_kernel_size,
mlp_ratio=mlp_ratio, backbone_expansion=backbone_expansion,
dropout=dropout, dropout=dropout,
append_observed_mask=True, append_observed_mask=True,
) )
@@ -440,9 +385,9 @@ class TimesNetExposureEncoder(nn.Module):
d_model=d_model, d_model=d_model,
n_layers=n_layers, n_layers=n_layers,
top_k=top_k, top_k=top_k,
n_convnext_blocks=n_convnext_blocks, n_backbone_blocks=n_backbone_blocks,
conv_kernel_size=conv_kernel_size, backbone_kernel_size=backbone_kernel_size,
mlp_ratio=mlp_ratio, backbone_expansion=backbone_expansion,
dropout=dropout, dropout=dropout,
append_observed_mask=True, append_observed_mask=True,
) )
@@ -485,10 +430,10 @@ class TimesNetSequenceDecoder(nn.Module):
latent_dim: int, latent_dim: int,
d_model: int, d_model: int,
n_layers: int = 2, n_layers: int = 2,
top_k: int = 3, top_k: int = 2,
n_convnext_blocks: int = 2, n_backbone_blocks: int = 1,
conv_kernel_size: int = 7, backbone_kernel_size: int = 5,
mlp_ratio: float = 4.0, backbone_expansion: float = 2.0,
dropout: float = 0.0, dropout: float = 0.0,
): ):
super().__init__() super().__init__()
@@ -498,9 +443,9 @@ class TimesNetSequenceDecoder(nn.Module):
TimesNetBlock( TimesNetBlock(
d_model=d_model, d_model=d_model,
top_k=top_k, top_k=top_k,
n_convnext_blocks=n_convnext_blocks, n_backbone_blocks=n_backbone_blocks,
conv_kernel_size=conv_kernel_size, backbone_kernel_size=backbone_kernel_size,
mlp_ratio=mlp_ratio, backbone_expansion=backbone_expansion,
dropout=dropout, dropout=dropout,
) )
for _ in range(n_layers) for _ in range(n_layers)
@@ -538,10 +483,10 @@ class TimesNetExposureAutoencoder(nn.Module):
monthly_input_dim: int = 2, monthly_input_dim: int = 2,
d_model: int | None = None, d_model: int | None = None,
n_layers: int = 2, n_layers: int = 2,
top_k: int = 3, top_k: int = 2,
n_convnext_blocks: int = 2, n_backbone_blocks: int = 1,
conv_kernel_size: int = 7, backbone_kernel_size: int = 5,
mlp_ratio: float = 4.0, backbone_expansion: float = 2.0,
dropout: float = 0.0, dropout: float = 0.0,
): ):
super().__init__() super().__init__()
@@ -553,9 +498,9 @@ class TimesNetExposureAutoencoder(nn.Module):
d_model=d_model, d_model=d_model,
n_layers=n_layers, n_layers=n_layers,
top_k=top_k, top_k=top_k,
n_convnext_blocks=n_convnext_blocks, n_backbone_blocks=n_backbone_blocks,
conv_kernel_size=conv_kernel_size, backbone_kernel_size=backbone_kernel_size,
mlp_ratio=mlp_ratio, backbone_expansion=backbone_expansion,
dropout=dropout, dropout=dropout,
use_gate=True, use_gate=True,
) )
@@ -564,9 +509,9 @@ class TimesNetExposureAutoencoder(nn.Module):
d_model=d_model, d_model=d_model,
n_layers=n_layers, n_layers=n_layers,
top_k=top_k, top_k=top_k,
n_convnext_blocks=n_convnext_blocks, n_backbone_blocks=n_backbone_blocks,
conv_kernel_size=conv_kernel_size, backbone_kernel_size=backbone_kernel_size,
mlp_ratio=mlp_ratio, backbone_expansion=backbone_expansion,
dropout=dropout, dropout=dropout,
) )
self.encoder = TimesNetExposureEncoder(**encoder_kwargs) self.encoder = TimesNetExposureEncoder(**encoder_kwargs)

10
evaluate_auc.py
View File

@@ -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")), dist_mode=str(cfg_get(args, cfg, "dist_mode", "exponential")),
dropout=float(cfg_get(args, cfg, "dropout", 0.0)), dropout=float(cfg_get(args, cfg, "dropout", 0.0)),
use_exposure_encoder=bool(cfg_get(args, cfg, "use_exposure_encoder", False)), 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_n_layers=int(cfg_get(args, cfg, "exposure_n_layers", 2)),
exposure_top_k=int(cfg_get(args, cfg, "exposure_top_k", 3)), exposure_top_k=int(cfg_get(args, cfg, "exposure_top_k", 2)),
exposure_n_convnext_blocks=int(cfg_get(args, cfg, "exposure_n_convnext_blocks", 2)), exposure_n_backbone_blocks=int(cfg_get(args, cfg, "exposure_n_backbone_blocks", 1)),
exposure_conv_kernel_size=int(cfg_get(args, cfg, "exposure_conv_kernel_size", 7)), exposure_backbone_kernel_size=int(cfg_get(args, cfg, "exposure_backbone_kernel_size", 5)),
exposure_mlp_ratio=float(cfg_get(args, cfg, "exposure_mlp_ratio", 4.0)), 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)), exposure_use_gate=bool(cfg_get(args, cfg, "exposure_use_gate", True)),
) )

16
models.py
View File

@@ -33,12 +33,12 @@ class DeepHealth(nn.Module):
use_exposure_encoder: bool = False, use_exposure_encoder: bool = False,
exposure_daily_input_dim: int = 4, exposure_daily_input_dim: int = 4,
exposure_monthly_input_dim: int = 2, exposure_monthly_input_dim: int = 2,
exposure_d_model: int | None = None, exposure_d_model: int | None = 64,
exposure_n_layers: int = 2, exposure_n_layers: int = 2,
exposure_top_k: int = 3, exposure_top_k: int = 2,
exposure_n_convnext_blocks: int = 2, exposure_n_backbone_blocks: int = 1,
exposure_conv_kernel_size: int = 7, exposure_backbone_kernel_size: int = 5,
exposure_mlp_ratio: float = 4.0, exposure_backbone_expansion: float = 2.0,
exposure_use_gate: bool = True, exposure_use_gate: bool = True,
): ):
super().__init__() super().__init__()
@@ -64,9 +64,9 @@ class DeepHealth(nn.Module):
d_model=exposure_d_model, d_model=exposure_d_model,
n_layers=exposure_n_layers, n_layers=exposure_n_layers,
top_k=exposure_top_k, top_k=exposure_top_k,
n_convnext_blocks=exposure_n_convnext_blocks, n_backbone_blocks=exposure_n_backbone_blocks,
conv_kernel_size=exposure_conv_kernel_size, backbone_kernel_size=exposure_backbone_kernel_size,
mlp_ratio=exposure_mlp_ratio, backbone_expansion=exposure_backbone_expansion,
dropout=dropout, dropout=dropout,
use_gate=exposure_use_gate, use_gate=exposure_use_gate,
) )

23
train_exposure_autoencoder.py
View File

@@ -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 from __future__ import annotations
import argparse import argparse
@@ -46,7 +46,7 @@ class ExposureWindowDataset(Dataset):
def parse_args() -> argparse.Namespace: def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser( 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("--exposure_cache_dir", required=True)
parser.add_argument("--train_eid_file", default="ukb_train_eid.csv") 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("--runs_root", default="runs")
parser.add_argument("--seed", type=int, default=42) parser.add_argument("--seed", type=int, default=42)
parser.add_argument("--n_embd", type=int, default=120) 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("--n_layers", type=int, default=2)
parser.add_argument("--top_k", type=int, default=3) parser.add_argument("--top_k", type=int, default=2)
parser.add_argument("--n_convnext_blocks", type=int, default=2) parser.add_argument("--n_backbone_blocks", type=int, default=1)
parser.add_argument("--conv_kernel_size", type=int, default=7) parser.add_argument("--backbone_kernel_size", type=int, default=5)
parser.add_argument("--mlp_ratio", type=float, default=4.0) parser.add_argument("--backbone_expansion", type=float, default=2.0)
parser.add_argument("--dropout", type=float, default=0.0) parser.add_argument("--dropout", type=float, default=0.0)
parser.add_argument("--mask_ratio", type=float, default=0.25) parser.add_argument("--mask_ratio", type=float, default=0.25)
parser.add_argument("--batch_size", type=int, default=16) parser.add_argument("--batch_size", type=int, default=16)
@@ -288,8 +288,9 @@ def main() -> None:
) )
model = TimesNetExposureAutoencoder( model = TimesNetExposureAutoencoder(
n_embd=args.n_embd, d_model=args.d_model, n_layers=args.n_layers, 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, top_k=args.top_k, n_backbone_blocks=args.n_backbone_blocks,
conv_kernel_size=args.conv_kernel_size, mlp_ratio=args.mlp_ratio, backbone_kernel_size=args.backbone_kernel_size,
backbone_expansion=args.backbone_expansion,
dropout=args.dropout, dropout=args.dropout,
).to(device) ).to(device)
model = maybe_wrap_data_parallel(model, args, device, logger) model = maybe_wrap_data_parallel(model, args, device, logger)
@@ -349,8 +350,8 @@ def main() -> None:
"model_config": { "model_config": {
key: config[key] for key in ( key: config[key] for key in (
"n_embd", "d_model", "n_layers", "top_k", "n_embd", "d_model", "n_layers", "top_k",
"n_convnext_blocks", "conv_kernel_size", "n_backbone_blocks", "backbone_kernel_size",
"mlp_ratio", "dropout", "backbone_expansion", "dropout",
) )
}, },
"normalization": { "normalization": {

22
train_next_step.py
View File

@@ -129,12 +129,12 @@ def parse_args() -> argparse.Namespace:
parser.add_argument("--dropout", type=float, default=0.0) parser.add_argument("--dropout", type=float, default=0.0)
parser.add_argument("--exposure_cache_dir", type=str, default=None) parser.add_argument("--exposure_cache_dir", type=str, default=None)
parser.add_argument("--mask_onset_exposure", action="store_true") 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_n_layers", type=int, default=2)
parser.add_argument("--exposure_top_k", type=int, default=3) parser.add_argument("--exposure_top_k", type=int, default=2)
parser.add_argument("--exposure_n_convnext_blocks", type=int, default=2) parser.add_argument("--exposure_n_backbone_blocks", type=int, default=1)
parser.add_argument("--exposure_conv_kernel_size", type=int, default=7) parser.add_argument("--exposure_backbone_kernel_size", type=int, default=5)
parser.add_argument("--exposure_mlp_ratio", type=float, default=4.0) parser.add_argument("--exposure_backbone_expansion", type=float, default=2.0)
parser.add_argument("--no_exposure_gate", action="store_true") parser.add_argument("--no_exposure_gate", action="store_true")
parser.add_argument("--target_mode", type=str, default="uts", parser.add_argument("--target_mode", type=str, default="uts",
choices=["delphi2m", "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_d_model=args.exposure_d_model,
exposure_n_layers=args.exposure_n_layers, exposure_n_layers=args.exposure_n_layers,
exposure_top_k=args.exposure_top_k, exposure_top_k=args.exposure_top_k,
exposure_n_convnext_blocks=args.exposure_n_convnext_blocks, exposure_n_backbone_blocks=args.exposure_n_backbone_blocks,
exposure_conv_kernel_size=args.exposure_conv_kernel_size, exposure_backbone_kernel_size=args.exposure_backbone_kernel_size,
exposure_mlp_ratio=args.exposure_mlp_ratio, exposure_backbone_expansion=args.exposure_backbone_expansion,
exposure_use_gate=not args.no_exposure_gate, exposure_use_gate=not args.no_exposure_gate,
) )
@@ -490,9 +490,9 @@ def build_metadata(
"exposure_d_model": args.exposure_d_model, "exposure_d_model": args.exposure_d_model,
"exposure_n_layers": int(args.exposure_n_layers), "exposure_n_layers": int(args.exposure_n_layers),
"exposure_top_k": int(args.exposure_top_k), "exposure_top_k": int(args.exposure_top_k),
"exposure_n_convnext_blocks": int(args.exposure_n_convnext_blocks), "exposure_n_backbone_blocks": int(args.exposure_n_backbone_blocks),
"exposure_conv_kernel_size": int(args.exposure_conv_kernel_size), "exposure_backbone_kernel_size": int(args.exposure_backbone_kernel_size),
"exposure_mlp_ratio": float(args.exposure_mlp_ratio), "exposure_backbone_expansion": float(args.exposure_backbone_expansion),
"exposure_use_gate": not bool(args.no_exposure_gate), "exposure_use_gate": not bool(args.no_exposure_gate),
"num_workers": int(args.num_workers), "num_workers": int(args.num_workers),
"prefetch_factor": int(args.prefetch_factor), "prefetch_factor": int(args.prefetch_factor),