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:
327
src_new/train.py
Normal file
327
src_new/train.py
Normal file
@@ -0,0 +1,327 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import random
|
||||
from dataclasses import asdict
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import pandas as pd
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
from torch.optim import AdamW
|
||||
from torch.optim.lr_scheduler import ReduceLROnPlateau
|
||||
|
||||
try:
|
||||
from .config import ExperimentConfig, make_rms_forward_config
|
||||
from .dataset import build_dataloaders, report_to_text
|
||||
from .model import build_model, compute_receptive_field
|
||||
except ImportError:
|
||||
from config import ExperimentConfig, make_rms_forward_config
|
||||
from dataset import build_dataloaders, report_to_text
|
||||
from model import build_model, compute_receptive_field
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="Train TCN for direct RMS regression.")
|
||||
parser.add_argument("--epochs", type=int, default=None)
|
||||
parser.add_argument("--batch-size", type=int, default=None)
|
||||
parser.add_argument("--device", type=str, default=None)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def set_seed(seed: int) -> None:
|
||||
random.seed(seed)
|
||||
torch.manual_seed(seed)
|
||||
torch.cuda.manual_seed_all(seed)
|
||||
|
||||
|
||||
def resolve_device(device_name: str) -> torch.device:
|
||||
if device_name.startswith("cuda") and not torch.cuda.is_available():
|
||||
return torch.device("cpu")
|
||||
return torch.device(device_name)
|
||||
|
||||
|
||||
def serialize_for_checkpoint(value: Any) -> Any:
|
||||
if isinstance(value, Path):
|
||||
return str(value)
|
||||
if isinstance(value, dict):
|
||||
return {key: serialize_for_checkpoint(sub_value) for key, sub_value in value.items()}
|
||||
if isinstance(value, tuple):
|
||||
return [serialize_for_checkpoint(item) for item in value]
|
||||
if isinstance(value, list):
|
||||
return [serialize_for_checkpoint(item) for item in value]
|
||||
return value
|
||||
|
||||
|
||||
def denormalize_target(
|
||||
pred_norm: torch.Tensor,
|
||||
target_norm: torch.Tensor,
|
||||
target_mean: torch.Tensor,
|
||||
target_std: torch.Tensor,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
pred = pred_norm * target_std + target_mean
|
||||
target = target_norm * target_std + target_mean
|
||||
return pred, target
|
||||
|
||||
|
||||
def relative_rms_error(pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
|
||||
return torch.abs(pred - target) / torch.clamp(target.abs(), min=1e-6)
|
||||
|
||||
|
||||
class RMSRegressionLoss(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
relative_rms_weight: float,
|
||||
log_rms_weight: float,
|
||||
mae_weight: float,
|
||||
waveform_l1_weight: float,
|
||||
waveform_huber_weight: float,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.relative_rms_weight = relative_rms_weight
|
||||
self.log_rms_weight = log_rms_weight
|
||||
self.mae_weight = mae_weight
|
||||
self.waveform_l1_weight = waveform_l1_weight
|
||||
self.waveform_huber_weight = waveform_huber_weight
|
||||
|
||||
def forward(
|
||||
self,
|
||||
pred: torch.Tensor,
|
||||
target: torch.Tensor,
|
||||
pred_wave: torch.Tensor,
|
||||
target_wave: torch.Tensor,
|
||||
mask: torch.Tensor,
|
||||
) -> dict[str, torch.Tensor]:
|
||||
rel_loss = relative_rms_error(pred, target).mean()
|
||||
log_loss = F.huber_loss(torch.log(torch.clamp(pred, min=1e-6)), torch.log(torch.clamp(target, min=1e-6)))
|
||||
mae_loss = F.l1_loss(pred, target)
|
||||
mask_expanded = mask.unsqueeze(-1)
|
||||
valid_count = mask_expanded.sum().clamp_min(1.0)
|
||||
wave_residual = (pred_wave - target_wave) * mask_expanded
|
||||
waveform_l1 = torch.abs(wave_residual).sum() / valid_count
|
||||
waveform_huber = F.huber_loss(pred_wave * mask_expanded, target_wave * mask_expanded, reduction="sum") / valid_count
|
||||
total = (
|
||||
self.relative_rms_weight * rel_loss
|
||||
+ self.log_rms_weight * log_loss
|
||||
+ self.mae_weight * mae_loss
|
||||
+ self.waveform_l1_weight * waveform_l1
|
||||
+ self.waveform_huber_weight * waveform_huber
|
||||
)
|
||||
return {
|
||||
"total": total,
|
||||
"relative": rel_loss,
|
||||
"log": log_loss,
|
||||
"mae": mae_loss,
|
||||
"waveform_l1": waveform_l1,
|
||||
"waveform_huber": waveform_huber,
|
||||
}
|
||||
|
||||
|
||||
def run_epoch(
|
||||
model: nn.Module,
|
||||
dataloader: torch.utils.data.DataLoader,
|
||||
optimizer: AdamW | None,
|
||||
criterion: RMSRegressionLoss,
|
||||
target_mean: torch.Tensor,
|
||||
target_std: torch.Tensor,
|
||||
device: torch.device,
|
||||
grad_clip_norm: float,
|
||||
scaler: torch.cuda.amp.GradScaler,
|
||||
amp_enabled: bool,
|
||||
) -> dict[str, float]:
|
||||
is_train = optimizer is not None
|
||||
model.train(is_train)
|
||||
|
||||
total_loss_sum = 0.0
|
||||
relative_loss_sum = 0.0
|
||||
log_loss_sum = 0.0
|
||||
mae_loss_sum = 0.0
|
||||
waveform_l1_sum = 0.0
|
||||
waveform_huber_sum = 0.0
|
||||
rms_error_sum = 0.0
|
||||
sample_count = 0
|
||||
|
||||
for batch in dataloader:
|
||||
x = batch["x"].to(device)
|
||||
aux = batch["aux"].to(device)
|
||||
mask = batch["mask"].to(device)
|
||||
y_norm = batch["y"].to(device)
|
||||
y_wave = batch["y_wave"].to(device)
|
||||
x_rms_raw = batch["x_rms_raw"].to(device)
|
||||
y_rms_raw = batch["y_raw"].to(device)
|
||||
|
||||
if is_train:
|
||||
optimizer.zero_grad(set_to_none=True)
|
||||
|
||||
with torch.amp.autocast(device_type=device.type, enabled=amp_enabled):
|
||||
outputs = model(x, mask, aux)
|
||||
pred_target, _ = denormalize_target(outputs["rms"], y_norm, target_mean, target_std)
|
||||
pred = torch.exp(pred_target) * x_rms_raw
|
||||
target = y_rms_raw
|
||||
losses = criterion(pred, target, outputs["waveform"], y_wave, mask)
|
||||
|
||||
if is_train:
|
||||
scaler.scale(losses["total"]).backward()
|
||||
scaler.unscale_(optimizer)
|
||||
torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip_norm)
|
||||
scaler.step(optimizer)
|
||||
scaler.update()
|
||||
|
||||
batch_size = x.shape[0]
|
||||
total_loss_sum += losses["total"].detach().item() * batch_size
|
||||
relative_loss_sum += losses["relative"].detach().item() * batch_size
|
||||
log_loss_sum += losses["log"].detach().item() * batch_size
|
||||
mae_loss_sum += losses["mae"].detach().item() * batch_size
|
||||
waveform_l1_sum += losses["waveform_l1"].detach().item() * batch_size
|
||||
waveform_huber_sum += losses["waveform_huber"].detach().item() * batch_size
|
||||
rms_error_sum += relative_rms_error(pred.detach(), target.detach()).mean().item() * batch_size
|
||||
sample_count += batch_size
|
||||
|
||||
return {
|
||||
"loss": total_loss_sum / sample_count,
|
||||
"relative_loss": relative_loss_sum / sample_count,
|
||||
"log_loss": log_loss_sum / sample_count,
|
||||
"mae_loss": mae_loss_sum / sample_count,
|
||||
"waveform_l1": waveform_l1_sum / sample_count,
|
||||
"waveform_huber": waveform_huber_sum / sample_count,
|
||||
"rms_error": rms_error_sum / sample_count,
|
||||
}
|
||||
|
||||
|
||||
def checkpoint_paths(config: ExperimentConfig) -> tuple[Path, Path]:
|
||||
root = config.data.project_root / config.train.checkpoint_dir / "forward_rms"
|
||||
root.mkdir(parents=True, exist_ok=True)
|
||||
return root / config.train.best_model_name, root / config.train.history_name
|
||||
|
||||
def train_model(config: ExperimentConfig) -> None:
|
||||
set_seed(config.train.seed)
|
||||
device = resolve_device(config.train.device)
|
||||
loaders, datasets, reports = build_dataloaders(config)
|
||||
train_dataset = datasets["train"]
|
||||
normalization = train_dataset.normalization
|
||||
target_mean = normalization.target_mean.to(device).view(1, 1)
|
||||
target_std = normalization.target_std.to(device).view(1, 1)
|
||||
|
||||
model = build_model(config).to(device)
|
||||
receptive_field = compute_receptive_field(config.model)
|
||||
optimizer = AdamW(model.parameters(), lr=config.train.learning_rate, weight_decay=config.train.weight_decay)
|
||||
scheduler = ReduceLROnPlateau(
|
||||
optimizer,
|
||||
mode="min",
|
||||
factor=config.train.lr_scheduler_factor,
|
||||
patience=config.train.lr_scheduler_patience,
|
||||
min_lr=config.train.min_learning_rate,
|
||||
)
|
||||
criterion = RMSRegressionLoss(
|
||||
relative_rms_weight=config.loss.relative_rms_weight,
|
||||
log_rms_weight=config.loss.log_rms_weight,
|
||||
mae_weight=config.loss.mae_weight,
|
||||
waveform_l1_weight=config.loss.waveform_l1_weight,
|
||||
waveform_huber_weight=config.loss.waveform_huber_weight,
|
||||
)
|
||||
amp_enabled = config.train.use_amp and device.type == "cuda"
|
||||
scaler = torch.cuda.amp.GradScaler(enabled=amp_enabled)
|
||||
|
||||
best_val_error = float("inf")
|
||||
epochs_without_improvement = 0
|
||||
history: list[dict[str, float]] = []
|
||||
best_model_path, history_path = checkpoint_paths(config)
|
||||
|
||||
print(f"Device: {device}")
|
||||
print(f"TCN receptive field: {receptive_field.receptive_field} samples, dilations={receptive_field.dilations}")
|
||||
print(report_to_text(reports))
|
||||
|
||||
for epoch in range(1, config.train.epochs + 1):
|
||||
train_metrics = run_epoch(
|
||||
model=model,
|
||||
dataloader=loaders["train"],
|
||||
optimizer=optimizer,
|
||||
criterion=criterion,
|
||||
target_mean=target_mean,
|
||||
target_std=target_std,
|
||||
device=device,
|
||||
grad_clip_norm=config.train.grad_clip_norm,
|
||||
scaler=scaler,
|
||||
amp_enabled=amp_enabled,
|
||||
)
|
||||
val_metrics = run_epoch(
|
||||
model=model,
|
||||
dataloader=loaders["val"],
|
||||
optimizer=None,
|
||||
criterion=criterion,
|
||||
target_mean=target_mean,
|
||||
target_std=target_std,
|
||||
device=device,
|
||||
grad_clip_norm=config.train.grad_clip_norm,
|
||||
scaler=scaler,
|
||||
amp_enabled=amp_enabled,
|
||||
)
|
||||
scheduler.step(val_metrics["rms_error"])
|
||||
current_lr = optimizer.param_groups[0]["lr"]
|
||||
|
||||
history_row = {
|
||||
"epoch": epoch,
|
||||
"lr": current_lr,
|
||||
"train_loss": train_metrics["loss"],
|
||||
"train_relative_loss": train_metrics["relative_loss"],
|
||||
"train_log_loss": train_metrics["log_loss"],
|
||||
"train_mae_loss": train_metrics["mae_loss"],
|
||||
"train_waveform_l1": train_metrics["waveform_l1"],
|
||||
"train_waveform_huber": train_metrics["waveform_huber"],
|
||||
"train_rms_error": train_metrics["rms_error"],
|
||||
"val_loss": val_metrics["loss"],
|
||||
"val_relative_loss": val_metrics["relative_loss"],
|
||||
"val_log_loss": val_metrics["log_loss"],
|
||||
"val_mae_loss": val_metrics["mae_loss"],
|
||||
"val_waveform_l1": val_metrics["waveform_l1"],
|
||||
"val_waveform_huber": val_metrics["waveform_huber"],
|
||||
"val_rms_error": val_metrics["rms_error"],
|
||||
}
|
||||
history.append(history_row)
|
||||
|
||||
print(
|
||||
f"Epoch {epoch:03d} | train_loss={train_metrics['loss']:.6f} | "
|
||||
f"val_loss={val_metrics['loss']:.6f} | val_rms_error={val_metrics['rms_error']:.6f} | lr={current_lr:.2e}"
|
||||
)
|
||||
|
||||
if val_metrics["rms_error"] < best_val_error:
|
||||
best_val_error = val_metrics["rms_error"]
|
||||
epochs_without_improvement = 0
|
||||
torch.save(
|
||||
{
|
||||
"epoch": epoch,
|
||||
"model_state_dict": model.state_dict(),
|
||||
"optimizer_state_dict": optimizer.state_dict(),
|
||||
"best_val_rms_error": best_val_error,
|
||||
"config": serialize_for_checkpoint(asdict(config)),
|
||||
},
|
||||
best_model_path,
|
||||
)
|
||||
else:
|
||||
epochs_without_improvement += 1
|
||||
if epochs_without_improvement >= config.train.early_stop_patience:
|
||||
print(f"Early stopping triggered after {epoch} epochs.")
|
||||
break
|
||||
|
||||
history_df = pd.DataFrame(history)
|
||||
history_df.to_csv(history_path, index=False)
|
||||
print(f"Best model saved to: {best_model_path}")
|
||||
print(f"Training history saved to: {history_path}")
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
config = make_rms_forward_config()
|
||||
if args.epochs is not None:
|
||||
config.train.epochs = args.epochs
|
||||
if args.batch_size is not None:
|
||||
config.data.batch_size = args.batch_size
|
||||
if args.device is not None:
|
||||
config.train.device = args.device
|
||||
train_model(config)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user