modified: Figure_1.png
modified: checkpoints/forward/best_tcn_model.pt modified: checkpoints/forward/training_history.csv new file: checkpoints_mlp/task1_feature_mlp/best_feature_mlp.pt new file: checkpoints_mlp/task1_feature_mlp/training_history.csv new file: checkpoints_rms/forward_rms/best_rms_model.pt new file: checkpoints_rms/forward_rms/training_history.csv modified: evaluation_outputs/forward/evaluation_forward_test_b0_s0.png new file: evaluation_outputs/forward/evaluation_forward_val_b0_s0.png new file: evaluation_outputs/forward/evaluation_forward_val_b3_s0.png new file: evaluation_outputs/forward_rms/evaluation_test_all_samples.csv new file: evaluation_outputs/forward_rms/evaluation_test_s0.png new file: evaluation_outputs/forward_rms/evaluation_train_all_samples.csv new file: evaluation_outputs/forward_rms/evaluation_val_all_samples.csv new file: evaluation_outputs/forward_rms/evaluation_val_s0.png new file: evaluation_outputs/forward_rms/evaluation_val_s0_waveform.png new file: evaluation_outputs/task1_feature_mlp/evaluation_train_all_samples.csv new file: evaluation_outputs/task1_feature_mlp/evaluation_train_curve.png new file: evaluation_outputs/task1_feature_mlp/evaluation_val_all_samples.csv new file: evaluation_outputs/task1_feature_mlp/evaluation_val_curve.png new file: evaluation_outputs/task1_feature_mlp/evaluation_val_s0.png new file: evaluation_outputs/task1_feature_mlp/harmonic_5mm_0.75Hz_prediction.png new file: evaluation_outputs/task1_feature_mlp/harmonic_5mm_1.55Hz_prediction.png new file: scripts/__pycache__/config.cpython-310.pyc new file: scripts/__pycache__/dataset.cpython-310.pyc new file: scripts/__pycache__/model.cpython-310.pyc new file: scripts/config.py new file: scripts/dataset.py new file: scripts/evaluate.py new file: scripts/model.py new file: scripts/predict_single.py new file: scripts/train.py modified: src/__pycache__/config.cpython-310.pyc modified: src/__pycache__/dataset.cpython-310.pyc modified: src/__pycache__/model.cpython-310.pyc modified: src/config.py modified: src/dataset.py modified: src/model.py new file: src_new/__pycache__/config.cpython-310.pyc new file: src_new/__pycache__/dataset.cpython-310.pyc new file: src_new/__pycache__/evaluate.cpython-310.pyc new file: src_new/__pycache__/model.cpython-310.pyc new file: src_new/__pycache__/train.cpython-310.pyc new file: src_new/config.py new file: src_new/dataset.py new file: src_new/evaluate.py new file: src_new/model.py new file: src_new/train.py
This commit is contained in:
140
src_new/model.py
Normal file
140
src_new/model.py
Normal file
@@ -0,0 +1,140 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.nn.utils import weight_norm
|
||||
|
||||
try:
|
||||
from .config import ExperimentConfig, ModelConfig
|
||||
except ImportError:
|
||||
from config import ExperimentConfig, ModelConfig
|
||||
|
||||
|
||||
class Chomp1d(nn.Module):
|
||||
def __init__(self, chomp_size: int) -> None:
|
||||
super().__init__()
|
||||
self.chomp_size = chomp_size
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
if self.chomp_size == 0:
|
||||
return x
|
||||
return x[:, :, :-self.chomp_size].contiguous()
|
||||
|
||||
|
||||
class CausalConv1d(nn.Module):
|
||||
def __init__(self, in_channels: int, out_channels: int, kernel_size: int, dilation: int = 1) -> None:
|
||||
super().__init__()
|
||||
padding = (kernel_size - 1) * dilation
|
||||
self.net = nn.Sequential(
|
||||
weight_norm(
|
||||
nn.Conv1d(
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
kernel_size=kernel_size,
|
||||
padding=padding,
|
||||
dilation=dilation,
|
||||
)
|
||||
),
|
||||
Chomp1d(padding),
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return self.net(x)
|
||||
|
||||
|
||||
class TemporalBlock(nn.Module):
|
||||
def __init__(self, in_channels: int, out_channels: int, kernel_size: int, dilation: int, dropout: float) -> None:
|
||||
super().__init__()
|
||||
self.conv1 = CausalConv1d(in_channels, out_channels, kernel_size, dilation=dilation)
|
||||
self.act1 = nn.GELU()
|
||||
self.dropout1 = nn.Dropout(dropout)
|
||||
self.conv2 = CausalConv1d(out_channels, out_channels, kernel_size, dilation=dilation)
|
||||
self.act2 = nn.GELU()
|
||||
self.dropout2 = nn.Dropout(dropout)
|
||||
self.residual = nn.Conv1d(in_channels, out_channels, kernel_size=1) if in_channels != out_channels else nn.Identity()
|
||||
self.final_act = nn.GELU()
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
residual = self.residual(x)
|
||||
out = self.dropout1(self.act1(self.conv1(x)))
|
||||
out = self.dropout2(self.act2(self.conv2(out)))
|
||||
return self.final_act(out + residual)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ReceptiveFieldInfo:
|
||||
receptive_field: int
|
||||
dilations: tuple[int, ...]
|
||||
|
||||
|
||||
class RMSRegressionTCN(nn.Module):
|
||||
def __init__(self, config: ModelConfig) -> None:
|
||||
super().__init__()
|
||||
blocks: list[nn.Module] = []
|
||||
in_channels = config.input_channels
|
||||
dilations: list[int] = []
|
||||
|
||||
for level, out_channels in enumerate(config.tcn_channels):
|
||||
dilation = config.dilation_base ** level
|
||||
dilations.append(dilation)
|
||||
blocks.append(
|
||||
TemporalBlock(
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
kernel_size=config.kernel_size,
|
||||
dilation=dilation,
|
||||
dropout=config.dropout,
|
||||
)
|
||||
)
|
||||
in_channels = out_channels
|
||||
|
||||
self.encoder = nn.Sequential(*blocks)
|
||||
self.rms_head = nn.Sequential(
|
||||
nn.Linear(in_channels * 2 + 3, config.pooled_feature_dim),
|
||||
nn.GELU(),
|
||||
nn.Dropout(config.dropout),
|
||||
nn.Linear(config.pooled_feature_dim, config.pooled_feature_dim),
|
||||
nn.GELU(),
|
||||
nn.Dropout(config.dropout),
|
||||
nn.Linear(config.pooled_feature_dim, 1),
|
||||
)
|
||||
self.waveform_head = nn.Sequential(
|
||||
nn.Conv1d(in_channels, in_channels, kernel_size=1),
|
||||
nn.GELU(),
|
||||
nn.Dropout(config.dropout),
|
||||
nn.Conv1d(in_channels, 1, kernel_size=1),
|
||||
)
|
||||
self.receptive_field_info = compute_receptive_field(config)
|
||||
|
||||
def forward(self, x: torch.Tensor, mask: torch.Tensor, aux: torch.Tensor) -> dict[str, torch.Tensor]:
|
||||
if x.ndim != 3:
|
||||
raise ValueError(f"Expected x shape (batch, seq, channels), got {tuple(x.shape)}")
|
||||
features = self.encoder(x.transpose(1, 2))
|
||||
mask_1d = mask.unsqueeze(1)
|
||||
masked_features = features * mask_1d
|
||||
valid_count = mask_1d.sum(dim=2).clamp_min(1.0)
|
||||
mean_pool = masked_features.sum(dim=2) / valid_count
|
||||
masked_for_max = features.masked_fill(mask_1d == 0.0, float("-inf"))
|
||||
max_pool = masked_for_max.max(dim=2).values
|
||||
max_pool = torch.where(torch.isfinite(max_pool), max_pool, torch.zeros_like(max_pool))
|
||||
fused = torch.cat([mean_pool, max_pool, aux], dim=1)
|
||||
rms_prediction = self.rms_head(fused)
|
||||
waveform_prediction = self.waveform_head(features).transpose(1, 2)
|
||||
return {"rms": rms_prediction, "waveform": waveform_prediction}
|
||||
|
||||
|
||||
def compute_receptive_field(config: ModelConfig) -> ReceptiveFieldInfo:
|
||||
receptive_field = 1
|
||||
dilations: list[int] = []
|
||||
for level, _ in enumerate(config.tcn_channels):
|
||||
dilation = config.dilation_base ** level
|
||||
dilations.append(dilation)
|
||||
receptive_field += 2 * (config.kernel_size - 1) * dilation
|
||||
return ReceptiveFieldInfo(receptive_field=receptive_field, dilations=tuple(dilations))
|
||||
|
||||
|
||||
def build_model(config: ExperimentConfig | ModelConfig) -> RMSRegressionTCN:
|
||||
model_config = config.model if isinstance(config, ExperimentConfig) else config
|
||||
return RMSRegressionTCN(model_config)
|
||||
Reference in New Issue
Block a user