"""Full 5000-epoch Burgers input-generalization experiment; validation selects best.""" import csv import os from pathlib import Path import time import torch from torch.utils.data import DataLoader, TensorDataset from common import (parser, load_config, seed_everything, device_for, load_splits, resolve, save_json, digest, build_model, checkpoint_save, evaluate, update_run_manifest) def main(): p = parser(__doc__) p.add_argument('--resume', default=None, help='Resume a locally generated last checkpoint') args = p.parse_args() cfg = load_config(args.config) tr = cfg['training'] if tr['epochs'] != 5000 or tr['batch_size'] != 32: raise ValueError('Formal experiment requires 5000 epochs and batch size 32') seed_everything(cfg['seed'], tr['num_threads']) device = device_for(tr['device']) result_dir = resolve(cfg['paths']['results']) result_dir.mkdir(parents=True, exist_ok=True) history_path = result_dir / 'training_history.csv' if history_path.exists() and not args.resume: raise FileExistsError('Training history already exists. Use --resume to continue this run.') splits, times, data_path = load_splits(cfg, ['train', 'validation']) train_u, train_y, train_ids = splits['train'] val_u, val_y, val_ids = splits['validation'] train_u, train_y = train_u.to(device), train_y.to(device) val_u, val_y = val_u.to(device), val_y.to(device) loader = DataLoader(TensorDataset(train_u, train_y), batch_size=32, shuffle=True, drop_last=False, num_workers=0) model = build_model(cfg, device) optimizer = torch.optim.Adam(model.parameters(), lr=tr['learning_rate'], weight_decay=tr['weight_decay']) scheduler = torch.optim.lr_scheduler.MultiplicativeLR( optimizer, lr_lambda=lambda _: tr['scheduler_gamma']) start_epoch, best_mse, best_epoch = 0, float('inf'), None if args.resume: old = torch.load(resolve(args.resume), map_location='cpu', weights_only=False) if old['model_config'] != cfg['model'] or old['data_sha256'] != digest(data_path): raise ValueError('Resume checkpoint architecture or source data differs') model.load_state_dict(old['model_state_dict']) optimizer.load_state_dict(old['optimizer_state_dict']) scheduler.load_state_dict(old['scheduler_state_dict']) torch.set_rng_state(old['torch_rng']) torch.cuda.set_rng_state_all(old['cuda_rng']) start_epoch, best_mse, best_epoch = old['epoch'], old['best_validation_mse'], old['best_epoch'] # Logs may contain epochs beyond the last atomic checkpoint after an # interrupted allocation. Replay those epochs from saved RNG/optimizer # state, keeping exactly one history row for every completed epoch. with history_path.open(newline='') as previous: reader = csv.DictReader(previous) fields = reader.fieldnames retained = [r for r in reader if int(r['epoch']) <= start_epoch] if [int(r['epoch']) for r in retained] != list(range(1, start_epoch + 1)): raise ValueError('Resume history is incomplete before the checkpoint epoch') temporary_history = history_path.with_suffix('.csv.tmp') with temporary_history.open('w', newline='') as previous: writer = csv.DictWriter(previous, fieldnames=fields) writer.writeheader() writer.writerows(retained) temporary_history.replace(history_path) print(f'Resuming after completed epoch {start_epoch}', flush=True) data_sha = digest(data_path) save_json(result_dir / 'run_manifest.json', { 'status': 'running', 'config': cfg, 'data_sha256': data_sha, 'slurm_job_id': os.environ.get('SLURM_JOB_ID'), 'torch_version': torch.__version__, 'train_indices': train_ids.tolist(), 'validation_indices': val_ids.tolist(), 'test_indices_reserved': [800, 1000], 'times': times.tolist(), 'model_parameters': sum(v.numel() for v in model.parameters()), 'test_used_for_selection': False, 'precision': 'float32'}) header = ['epoch', 'train_half_mse', 'train_mse', 'last_batch_half_mse', 'learning_rate', 'validation_mse', 'validation_rmse_mean', 'validation_rmse_std', 'elapsed_seconds'] started = time.monotonic() with history_path.open('a' if args.resume else 'w', newline='', buffering=1) as handle: writer = csv.DictWriter(handle, fieldnames=header) if not args.resume: writer.writeheader() for epoch_index in range(start_epoch, tr['epochs']): epoch = epoch_index + 1 model.train() total, seen, last_loss = 0., 0, None lr_used = optimizer.param_groups[0]['lr'] for u, target in loader: optimizer.zero_grad(set_to_none=True) predicted = model(u, times) loss = 0.5 * (predicted - target).square().mean() if not torch.isfinite(loss): raise RuntimeError(f'Nonfinite training loss at epoch {epoch}') loss.backward() optimizer.step() last_loss = float(loss.detach()) total += last_loss * len(u) seen += len(u) if seen != 600: raise RuntimeError(f'Only {seen} training samples processed') if epoch_index > 0 and epoch_index % tr['scheduler_frequency'] == 0: scheduler.step() row = {'epoch': epoch, 'train_half_mse': total / seen, 'train_mse': 2 * total / seen, 'last_batch_half_mse': last_loss, 'learning_rate': lr_used, 'elapsed_seconds': time.monotonic() - started} val_summary = None if epoch_index % tr['validation_frequency'] == 0 or epoch == tr['epochs']: val_summary, *_ = evaluate(model, val_u, val_y, times, tr['batch_size']) row.update(validation_mse=val_summary['mse'], validation_rmse_mean=val_summary['trajectory_rmse_mean'], validation_rmse_std=val_summary['trajectory_rmse_std_ddof1']) if val_summary['mse'] < best_mse: best_mse, best_epoch = val_summary['mse'], epoch checkpoint_save(cfg['paths']['checkpoint'], { 'format_version': 1, 'model_config': cfg['model'], 'model_state_dict': model.state_dict(), 'epoch': epoch, 'validation_mse': best_mse, 'validation_metrics': val_summary, 'selection_split': 'validation[600:800]', 'data_sha256': data_sha, 'seed': cfg['seed']}) save_json(result_dir / 'best_validation_metrics.json', {'epoch': epoch, 'split': 'validation', **val_summary}) print(f'VALIDATION epoch={epoch} mse={val_summary["mse"]:.9e} ' f'RMSE_mean={val_summary["trajectory_rmse_mean"]:.9e} ' f'RMSE_std={val_summary["trajectory_rmse_std_ddof1"]:.9e} ' f'best_epoch={best_epoch}', flush=True) writer.writerow(row) print(f'epoch={epoch:04d}/5000 train_loss(half_MSE)={total / seen:.9e} ' f'train_MSE={2 * total / seen:.9e} lr={lr_used:.6e} ' f'samples={seen} elapsed={row["elapsed_seconds"]:.1f}s', flush=True) if epoch % tr['checkpoint_frequency'] == 0 or epoch == tr['epochs']: checkpoint_save(cfg['paths']['last_checkpoint'], { 'model_config': cfg['model'], 'model_state_dict': model.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), 'scheduler_state_dict': scheduler.state_dict(), 'epoch': epoch, 'best_validation_mse': best_mse, 'best_epoch': best_epoch, 'torch_rng': torch.get_rng_state(), 'cuda_rng': torch.cuda.get_rng_state_all(), 'data_sha256': data_sha}) summary = {'status': 'completed', 'completed_epochs': tr['epochs'], 'batches_per_epoch': len(loader), 'training_samples_per_epoch': 600, 'best_epoch': best_epoch, 'best_validation_mse': best_mse, 'selection_split': 'validation[600:800]', 'test_used_for_selection': False, 'checkpoint': str(cfg['paths']['checkpoint']), 'data_sha256': data_sha, 'elapsed_seconds': time.monotonic() - started, 'slurm_job_id': os.environ.get('SLURM_JOB_ID')} save_json(result_dir / 'training_summary.json', summary) update_run_manifest(cfg, status='training_completed', completed_epochs=tr['epochs'], best_epoch=best_epoch, best_validation_mse=best_mse) print('FORMAL TRAINING COMPLETE', summary, flush=True) if __name__ == '__main__': main()