File size: 9,061 Bytes
e456924 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 | """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()
|