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
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)

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")),
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)),
)

16
models.py
View File

@@ -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,
)

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
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": {

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("--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),