| """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'] |
| |
| |
| |
| 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() |
|
|
