import argparse import os import random import gc import matplotlib.pyplot as plt import numpy as np import seaborn as sns sns.set_style('darkgrid') import torch if torch.cuda.is_available(): # For faster torch.set_float32_matmul_precision('high') import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP from torch.utils.data.distributed import DistributedSampler from utils.samplers import FastRandomDistributedSampler from tqdm.auto import tqdm # Data/Task Specific Imports from data.custom_datasets import MazeImageFolder # Model Imports from models.ctm import ContinuousThoughtMachine from models.lstm import LSTMBaseline from models.ff import FFBaseline # Plotting/Utils Imports from tasks.mazes.plotting import make_maze_gif from tasks.image_classification.plotting import plot_neural_dynamics from utils.housekeeping import set_seed, zip_python_code from utils.losses import maze_loss from utils.schedulers import WarmupCosineAnnealingLR, WarmupMultiStepLR, warmup import torchvision torchvision.disable_beta_transforms_warning() import warnings warnings.filterwarnings("ignore", message="using precomputed metric; inverse_transform will be unavailable") warnings.filterwarnings('ignore', message='divide by zero encountered in power', category=RuntimeWarning) warnings.filterwarnings( "ignore", "Corrupt EXIF data", UserWarning, r"^PIL\.TiffImagePlugin$" ) warnings.filterwarnings( "ignore", "UserWarning: Metadata Warning", UserWarning, r"^PIL\.TiffImagePlugin$" ) warnings.filterwarnings( "ignore", "UserWarning: Truncated File Read", UserWarning, r"^PIL\.TiffImagePlugin$" ) def parse_args(): parser = argparse.ArgumentParser() # Model Selection parser.add_argument('--model', type=str, required=True, choices=['ctm', 'lstm', 'ff'], help='Model type to train.') # Model Architecture parser.add_argument('--d_model', type=int, default=512, help='Dimension of the model.') parser.add_argument('--dropout', type=float, default=0.0, help='Dropout rate.') parser.add_argument('--backbone_type', type=str, default='resnet34-2', help='Type of backbone featureiser.') # CTM / LSTM specific parser.add_argument('--d_input', type=int, default=128, help='Dimension of the input (CTM, LSTM).') parser.add_argument('--heads', type=int, default=8, help='Number of attention heads (CTM, LSTM).') parser.add_argument('--iterations', type=int, default=75, help='Number of internal ticks (CTM, LSTM).') parser.add_argument('--positional_embedding_type', type=str, default='none', help='Type of positional embedding (CTM, LSTM).', choices=['none', 'learnable-fourier', 'multi-learnable-fourier', 'custom-rotational']) # CTM specific parser.add_argument('--synapse_depth', type=int, default=8, help='Depth of U-NET model for synapse. 1=linear, no unet (CTM only).') parser.add_argument('--n_synch_out', type=int, default=32, help='Number of neurons to use for output synch (CTM only).') parser.add_argument('--n_synch_action', type=int, default=32, help='Number of neurons to use for observation/action synch (CTM only).') parser.add_argument('--neuron_select_type', type=str, default='random-pairing', help='Protocol for selecting neuron subset (CTM only).') parser.add_argument('--n_random_pairing_self', type=int, default=0, help='Number of neurons paired self-to-self for synch (CTM only).') parser.add_argument('--memory_length', type=int, default=25, help='Length of the pre-activation history for NLMS (CTM only).') parser.add_argument('--deep_memory', action=argparse.BooleanOptionalAction, default=True, help='Use deep memory (CTM only).') parser.add_argument('--memory_hidden_dims', type=int, default=32, help='Hidden dimensions of the memory if using deep memory (CTM only).') parser.add_argument('--dropout_nlm', type=float, default=None, help='Dropout rate for NLMs specifically. Unset to match dropout on the rest of the model (CTM only).') parser.add_argument('--do_normalisation', action=argparse.BooleanOptionalAction, default=False, help='Apply normalization in NLMs (CTM only).') # LSTM specific parser.add_argument('--num_layers', type=int, default=2, help='Number of LSTM stacked layers (LSTM only).') # Task Specific Args parser.add_argument('--maze_route_length', type=int, default=100, help='Length to truncate targets.') parser.add_argument('--cirriculum_lookahead', type=int, default=5, help='How far to look ahead for cirriculum.') # Training parser.add_argument('--batch_size', type=int, default=16, help='Batch size for training (per GPU).') parser.add_argument('--batch_size_test', type=int, default=64, help='Batch size for testing (per GPU).') parser.add_argument('--lr', type=float, default=1e-4, help='Learning rate for the model.') parser.add_argument('--training_iterations', type=int, default=100001, help='Number of training iterations.') parser.add_argument('--warmup_steps', type=int, default=5000, help='Number of warmup steps.') parser.add_argument('--use_scheduler', action=argparse.BooleanOptionalAction, default=True, help='Use a learning rate scheduler.') parser.add_argument('--scheduler_type', type=str, default='cosine', choices=['multistep', 'cosine'], help='Type of learning rate scheduler.') parser.add_argument('--milestones', type=int, default=[8000, 15000, 20000], nargs='+', help='Learning rate scheduler milestones.') parser.add_argument('--gamma', type=float, default=0.1, help='Learning rate scheduler gamma for multistep.') parser.add_argument('--weight_decay', type=float, default=0.0, help='Weight decay factor.') parser.add_argument('--weight_decay_exclusion_list', type=str, nargs='+', default=[], help='List to exclude from weight decay. Typically good: bn, ln, bias, start') parser.add_argument('--num_workers_train', type=int, default=0, help='Num workers training.') parser.add_argument('--gradient_clipping', type=float, default=-1, help='Gradient quantile clipping value (-1 to disable).') parser.add_argument('--use_custom_sampler', action=argparse.BooleanOptionalAction, default=False, help='Use custom fast sampler to avoid reshuffling.') parser.add_argument('--do_compile', action=argparse.BooleanOptionalAction, default=False, help='Try to compile model components.') # Logging and Saving parser.add_argument('--log_dir', type=str, default='logs/scratch', help='Directory for logging.') parser.add_argument('--dataset', type=str, default='mazes-medium', help='Dataset to use.', choices=['mazes-medium', 'mazes-large']) parser.add_argument('--save_every', type=int, default=1000, help='Save checkpoints every this many iterations.') parser.add_argument('--seed', type=int, default=412, help='Random seed.') parser.add_argument('--reload', action=argparse.BooleanOptionalAction, default=False, help='Reload from disk?') parser.add_argument('--reload_model_only', action=argparse.BooleanOptionalAction, default=False, help='Reload only the model from disk?') # Default False based on user edit parser.add_argument('--strict_reload', action=argparse.BooleanOptionalAction, default=False, help='Should use strict reload for model weights.') parser.add_argument('--ignore_metrics_when_reloading', action=argparse.BooleanOptionalAction, default=False, help='Ignore metrics when reloading (for debugging)?') # Tracking parser.add_argument('--track_every', type=int, default=1000, help='Track metrics every this many iterations.') parser.add_argument('--n_test_batches', type=int, default=2, help='How many minibatches to approx metrics. Set to -1 for full eval') # Precision parser.add_argument('--use_amp', action=argparse.BooleanOptionalAction, default=False, help='AMP autocast.') args = parser.parse_args() return args # --- DDP Setup Functions --- def setup_ddp(): if 'RANK' not in os.environ: os.environ['RANK'] = '0' os.environ['WORLD_SIZE'] = '1' os.environ['MASTER_ADDR'] = 'localhost' os.environ['MASTER_PORT'] = '12356' # Different port from image classification os.environ['LOCAL_RANK'] = '0' print("Running in non-distributed mode (simulated DDP setup).") if not torch.cuda.is_available() or int(os.environ['WORLD_SIZE']) == 1: dist.init_process_group(backend='gloo') print("Initialized process group with Gloo backend for single/CPU process.") rank = int(os.environ['RANK']) world_size = int(os.environ['WORLD_SIZE']) local_rank = int(os.environ['LOCAL_RANK']) return rank, world_size, local_rank dist.init_process_group(backend='nccl') rank = int(os.environ['RANK']) world_size = int(os.environ['WORLD_SIZE']) local_rank = int(os.environ['LOCAL_RANK']) if torch.cuda.is_available(): torch.cuda.set_device(local_rank) print(f"Rank {rank} setup on GPU {local_rank}") else: print(f"Rank {rank} setup on CPU") return rank, world_size, local_rank def cleanup_ddp(): if dist.is_initialized(): dist.destroy_process_group() print("DDP cleanup complete.") def is_main_process(rank): return rank == 0 # --- End DDP Setup --- if __name__=='__main__': args = parse_args() rank, world_size, local_rank = setup_ddp() set_seed(args.seed + rank, False) # Rank 0 handles directory creation and initial logging if is_main_process(rank): if not os.path.exists(args.log_dir): os.makedirs(args.log_dir) zip_python_code(f'{args.log_dir}/repo_state.zip') with open(f'{args.log_dir}/args.txt', 'w') as f: print(args, file=f) if world_size > 1: dist.barrier() assert args.dataset in ['mazes-medium', 'mazes-large'] # Setup Device if torch.cuda.is_available(): device = torch.device(f'cuda:{local_rank}') else: device = torch.device('cpu') if world_size > 1: warnings.warn("Running DDP on CPU is not recommended.") if is_main_process(rank): print(f'Main process (Rank {rank}): Using device {device}. World size: {world_size}. Model: {args.model}') prediction_reshaper = [args.maze_route_length, 5] args.out_dims = args.maze_route_length * 5 # --- Model Definition (Conditional) --- model_base = None # Base model before DDP wrapping if args.model == 'ctm': model_base = ContinuousThoughtMachine( iterations=args.iterations, d_model=args.d_model, d_input=args.d_input, heads=args.heads, n_synch_out=args.n_synch_out, n_synch_action=args.n_synch_action, synapse_depth=args.synapse_depth, memory_length=args.memory_length, deep_nlms=args.deep_memory, memory_hidden_dims=args.memory_hidden_dims, do_layernorm_nlm=args.do_normalisation, backbone_type=args.backbone_type, positional_embedding_type=args.positional_embedding_type, out_dims=args.out_dims, prediction_reshaper=prediction_reshaper, dropout=args.dropout, dropout_nlm=args.dropout_nlm, neuron_select_type=args.neuron_select_type, n_random_pairing_self=args.n_random_pairing_self, ).to(device) elif args.model == 'lstm': model_base = LSTMBaseline( num_layers=args.num_layers, iterations=args.iterations, d_model=args.d_model, d_input=args.d_input, heads=args.heads, backbone_type=args.backbone_type, positional_embedding_type=args.positional_embedding_type, out_dims=args.out_dims, prediction_reshaper=prediction_reshaper, dropout=args.dropout, ).to(device) elif args.model == 'ff': model_base = FFBaseline( d_model=args.d_model, backbone_type=args.backbone_type, out_dims=args.out_dims, dropout=args.dropout, ).to(device) else: raise ValueError(f"Unknown model type: {args.model}") # Use pseudo-input *before* DDP wrapping try: # Determine pseudo input shape based on dataset h_w = 39 if args.dataset in ['mazes-small', 'mazes-medium'] else 99 # Example dimensions pseudo_inputs = torch.zeros((1, 3, h_w, h_w), device=device).float() model_base(pseudo_inputs) except Exception as e: print(f"Warning: Pseudo forward pass failed: {e}") if is_main_process(rank): print(f'Total params: {sum(p.numel() for p in model_base.parameters() if p.requires_grad)}') # Wrap model with DDP if device.type == 'cuda' and world_size > 1: model = DDP(model_base, device_ids=[local_rank], output_device=local_rank) elif device.type == 'cpu' and world_size > 1: model = DDP(model_base) else: model = model_base # --- End Model Definition --- # Data Loading (After model setup to allow pseudo pass first) dataset_mean = [0,0,0] dataset_std = [1,1,1] which_maze = args.dataset.split('-')[-1] data_root = f'data/mazes/{which_maze}' train_data = MazeImageFolder(root=f'{data_root}/train/', which_set='train', maze_route_length=args.maze_route_length) test_data = MazeImageFolder(root=f'{data_root}/test/', which_set='test', maze_route_length=args.maze_route_length) train_sampler = (FastRandomDistributedSampler(train_data, num_replicas=world_size, rank=rank, seed=args.seed, epoch_steps=int(10e10)) if args.use_custom_sampler else DistributedSampler(train_data, num_replicas=world_size, rank=rank, shuffle=True, seed=args.seed)) test_sampler = DistributedSampler(test_data, num_replicas=world_size, rank=rank, shuffle=False, seed=args.seed) num_workers_test = 1 trainloader = torch.utils.data.DataLoader(train_data, batch_size=args.batch_size, sampler=train_sampler, num_workers=args.num_workers_train, pin_memory=True, drop_last=True) testloader = torch.utils.data.DataLoader(test_data, batch_size=args.batch_size_test, sampler=test_sampler, num_workers=num_workers_test, pin_memory=True, drop_last=False) # Optimizer and scheduler decay_params = [] no_decay_params = [] no_decay_names = [] for name, param in model.named_parameters(): if not param.requires_grad: continue # Skip parameters that don't require gradients if any(exclusion_str in name for exclusion_str in args.weight_decay_exclusion_list): no_decay_params.append(param) no_decay_names.append(name) else: decay_params.append(param) if len(no_decay_names) and is_main_process(rank): print(f'WARNING, excluding: {no_decay_names}') # Optimizer and scheduler (Common setup) if len(no_decay_names) and args.weight_decay!=0: optimizer = torch.optim.AdamW([{'params': decay_params, 'weight_decay':args.weight_decay}, {'params': no_decay_params, 'weight_decay':0}], lr=args.lr, eps=1e-8 if not args.use_amp else 1e-6) else: optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, eps=1e-8 if not args.use_amp else 1e-6, weight_decay=args.weight_decay) warmup_schedule = warmup(args.warmup_steps) scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=warmup_schedule.step) if args.use_scheduler: if args.scheduler_type == 'multistep': scheduler = WarmupMultiStepLR(optimizer, warmup_steps=args.warmup_steps, milestones=args.milestones, gamma=args.gamma) elif args.scheduler_type == 'cosine': scheduler = WarmupCosineAnnealingLR(optimizer, args.warmup_steps, args.training_iterations, warmup_start_lr=1e-20, eta_min=1e-7) else: raise NotImplementedError # Metrics tracking (Rank 0 stores history) start_iter = 0 iters = [] train_losses, test_losses = [], [] train_accuracies, test_accuracies = [], [] # Avg Step Acc (scalar list) train_accuracies_most_certain, test_accuracies_most_certain = [], [] # Avg Step Acc @ Certain tick (scalar list) train_accuracies_most_certain_permaze, test_accuracies_most_certain_permaze = [], [] # Full Maze Acc @ Certain tick (scalar list) scaler = torch.amp.GradScaler("cuda" if device.type == 'cuda' else "cpu", enabled=args.use_amp) # Reloading Logic if args.reload: map_location = device chkpt_path = f'{args.log_dir}/checkpoint.pt' if os.path.isfile(chkpt_path): print(f'Rank {rank}: Reloading from: {chkpt_path}') if not args.strict_reload: print('WARNING: not using strict reload for model weights!') checkpoint = torch.load(chkpt_path, map_location=map_location, weights_only=False) model_to_load = model.module if isinstance(model, DDP) else model state_dict = checkpoint['model_state_dict'] has_module_prefix = all(k.startswith('module.') for k in state_dict) is_wrapped = isinstance(model, DDP) if has_module_prefix and not is_wrapped: state_dict = {k.partition('module.')[2]: v for k,v in state_dict.items()} elif not has_module_prefix and is_wrapped: load_result = model_to_load.load_state_dict(state_dict, strict=args.strict_reload) print(f" Loaded state_dict. Missing: {load_result.missing_keys}, Unexpected: {load_result.unexpected_keys}") state_dict = None # Prevent loading again if state_dict is not None: load_result = model_to_load.load_state_dict(state_dict, strict=args.strict_reload) print(f" Loaded state_dict. Missing: {load_result.missing_keys}, Unexpected: {load_result.unexpected_keys}") if not args.reload_model_only: print(f'Rank {rank}: Reloading optimizer, scheduler, scaler, iteration.') optimizer.load_state_dict(checkpoint['optimizer_state_dict']) scheduler.load_state_dict(checkpoint['scheduler_state_dict']) scaler.load_state_dict(checkpoint['scaler_state_dict']) start_iter = checkpoint['iteration'] if is_main_process(rank) and not args.ignore_metrics_when_reloading: print(f'Rank {rank}: Reloading metrics history.') iters = checkpoint['iters'] train_losses = checkpoint['train_losses'] test_losses = checkpoint['test_losses'] train_accuracies = checkpoint['train_accuracies'] # Reloading simplified avg step acc list test_accuracies = checkpoint['test_accuracies'] train_accuracies_most_certain = checkpoint['train_accuracies_most_certain'] test_accuracies_most_certain = checkpoint['test_accuracies_most_certain'] train_accuracies_most_certain_permaze = checkpoint['train_accuracies_most_certain_permaze'] test_accuracies_most_certain_permaze = checkpoint['test_accuracies_most_certain_permaze'] elif is_main_process(rank) and args.ignore_metrics_when_reloading: print(f'Rank {rank}: Ignoring metrics history upon reload.') else: print(f'Rank {rank}: Only reloading model weights!') if is_main_process(rank) and 'torch_rng_state' in checkpoint and not args.reload_model_only: print(f'Rank {rank}: Loading RNG states.') torch.set_rng_state(checkpoint['torch_rng_state'].cpu()) np.random.set_state(checkpoint['numpy_rng_state']) random.setstate(checkpoint['random_rng_state']) del checkpoint gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() print(f"Rank {rank}: Reload finished, starting from iteration {start_iter}") else: print(f"Rank {rank}: Checkpoint not found at {chkpt_path}, starting from scratch.") if world_size > 1: dist.barrier() # Conditional Compilation if args.do_compile: if is_main_process(rank): print('Compiling model components...') model_to_compile = model.module if isinstance(model, DDP) else model if hasattr(model_to_compile, 'backbone'): model_to_compile.backbone = torch.compile(model_to_compile.backbone, mode='reduce-overhead', fullgraph=True) if args.model == 'ctm': model_to_compile.synapses = torch.compile(model_to_compile.synapses, mode='reduce-overhead', fullgraph=True) if world_size > 1: dist.barrier() if is_main_process(rank): print('Compilation finished.') # --- Training Loop --- model.train() pbar = tqdm(total=args.training_iterations, initial=start_iter, leave=False, position=0, dynamic_ncols=True, disable=not is_main_process(rank)) iterator = iter(trainloader) for bi in range(start_iter, args.training_iterations): # --- Evaluation and Plotting (Rank 0 + Aggregation) --- if bi % args.track_every == 0 and (bi != 0 or args.reload_model_only): model.eval() with torch.inference_mode(): # --- Distributed Evaluation --- if is_main_process(rank): iters.append(bi) # Track iterations on rank 0 # Initialize accumulators on device total_train_loss = torch.tensor(0.0, device=device) total_train_correct_certain = torch.tensor(0.0, device=device) # Sum correct steps @ certain tick total_train_mazes_solved = torch.tensor(0.0, device=device) # Sum solved mazes @ certain tick total_train_steps = torch.tensor(0.0, device=device) # Total steps evaluated (B * S) total_train_mazes = torch.tensor(0.0, device=device) # Total mazes evaluated (B) # TRAIN METRICS train_eval_sampler = DistributedSampler(train_data, num_replicas=world_size, rank=rank, shuffle=False) train_eval_loader = torch.utils.data.DataLoader(train_data, batch_size=args.batch_size_test, sampler=train_eval_sampler, num_workers=num_workers_test, pin_memory=True) pbar_inner_desc = 'Eval Train (Rank 0)' if is_main_process(rank) else None with tqdm(total=len(train_eval_loader), desc=pbar_inner_desc, leave=False, position=1, dynamic_ncols=True, disable=not is_main_process(rank)) as pbar_inner: for inferi, (inputs, targets) in enumerate(train_eval_loader): inputs = inputs.to(device, non_blocking=True) targets = targets.to(device, non_blocking=True) # B, S batch_size = inputs.size(0) seq_len = targets.size(1) loss_eval = None pred_at_certain = None # Shape B, S if args.model == 'ctm': predictions_raw, certainties, _ = model(inputs) predictions = predictions_raw.reshape(batch_size, -1, 5, predictions_raw.size(-1)) # B,S,C,T loss_eval, where_most_certain, _ = maze_loss(predictions, certainties, targets, use_most_certain=True) pred_at_certain = predictions.argmax(2)[torch.arange(batch_size, device=device), :, where_most_certain] elif args.model == 'lstm': predictions_raw, certainties, _ = model(inputs) predictions = predictions_raw.reshape(batch_size, -1, 5, predictions_raw.size(-1)) # B,S,C,T loss_eval, where_most_certain, _ = maze_loss(predictions, certainties, targets, use_most_certain=False) # where = -1 pred_at_certain = predictions.argmax(2)[torch.arange(batch_size, device=device), :, where_most_certain] elif args.model == 'ff': predictions_raw = model(inputs) # B, S*C predictions = predictions_raw.reshape(batch_size, -1, 5) # B,S,C loss_eval, where_most_certain, _ = maze_loss(predictions.unsqueeze(-1), None, targets, use_most_certain=False) # where = -1 pred_at_certain = predictions.argmax(2) # Accumulate metrics total_train_loss += loss_eval * batch_size # Sum losses correct_steps = (pred_at_certain == targets) # B, S boolean total_train_correct_certain += correct_steps.sum() # Sum correct steps across batch total_train_mazes_solved += correct_steps.all(dim=-1).sum() # Sum mazes where all steps are correct total_train_steps += batch_size * seq_len total_train_mazes += batch_size if args.n_test_batches != -1 and inferi >= args.n_test_batches -1: break pbar_inner.update(1) # Aggregate Train Metrics if world_size > 1: dist.all_reduce(total_train_loss, op=dist.ReduceOp.SUM) dist.all_reduce(total_train_correct_certain, op=dist.ReduceOp.SUM) dist.all_reduce(total_train_mazes_solved, op=dist.ReduceOp.SUM) dist.all_reduce(total_train_steps, op=dist.ReduceOp.SUM) dist.all_reduce(total_train_mazes, op=dist.ReduceOp.SUM) # Calculate final Train metrics on Rank 0 if is_main_process(rank) and total_train_mazes > 0: avg_train_loss = total_train_loss.item() / total_train_mazes.item() # Avg loss per maze/sample avg_train_acc_step = total_train_correct_certain.item() / total_train_steps.item() # Avg correct step % avg_train_acc_maze = total_train_mazes_solved.item() / total_train_mazes.item() # Avg full maze solved % train_losses.append(avg_train_loss) train_accuracies_most_certain.append(avg_train_acc_step) train_accuracies_most_certain_permaze.append(avg_train_acc_maze) # train_accuracies list remains unused/placeholder for this simplified metric structure print(f"Iter {bi} Train Metrics (Agg): Loss={avg_train_loss:.4f}, StepAcc={avg_train_acc_step:.4f}, MazeAcc={avg_train_acc_maze:.4f}") # TEST METRICS total_test_loss = torch.tensor(0.0, device=device) total_test_correct_certain = torch.tensor(0.0, device=device) total_test_mazes_solved = torch.tensor(0.0, device=device) total_test_steps = torch.tensor(0.0, device=device) total_test_mazes = torch.tensor(0.0, device=device) pbar_inner_desc = 'Eval Test (Rank 0)' if is_main_process(rank) else None with tqdm(total=len(testloader), desc=pbar_inner_desc, leave=False, position=1, dynamic_ncols=True, disable=not is_main_process(rank)) as pbar_inner: for inferi, (inputs, targets) in enumerate(testloader): inputs = inputs.to(device, non_blocking=True) targets = targets.to(device, non_blocking=True) batch_size = inputs.size(0) seq_len = targets.size(1) loss_eval = None pred_at_certain = None if args.model == 'ctm': predictions_raw, certainties, _ = model(inputs) predictions = predictions_raw.reshape(batch_size, -1, 5, predictions_raw.size(-1)) loss_eval, where_most_certain, _ = maze_loss(predictions, certainties, targets, use_most_certain=True) pred_at_certain = predictions.argmax(2)[torch.arange(batch_size, device=device), :, where_most_certain] elif args.model == 'lstm': predictions_raw, certainties, _ = model(inputs) predictions = predictions_raw.reshape(batch_size, -1, 5, predictions_raw.size(-1)) loss_eval, where_most_certain, _ = maze_loss(predictions, certainties, targets, use_most_certain=False) pred_at_certain = predictions.argmax(2)[torch.arange(batch_size, device=device), :, where_most_certain] elif args.model == 'ff': predictions_raw = model(inputs) predictions = predictions_raw.reshape(batch_size, -1, 5) loss_eval, where_most_certain, _ = maze_loss(predictions.unsqueeze(-1), None, targets, use_most_certain=False) pred_at_certain = predictions.argmax(2) total_test_loss += loss_eval * batch_size correct_steps = (pred_at_certain == targets) total_test_correct_certain += correct_steps.sum() total_test_mazes_solved += correct_steps.all(dim=-1).sum() total_test_steps += batch_size * seq_len total_test_mazes += batch_size if args.n_test_batches != -1 and inferi >= args.n_test_batches -1: break pbar_inner.update(1) # Aggregate Test Metrics if world_size > 1: dist.all_reduce(total_test_loss, op=dist.ReduceOp.SUM) dist.all_reduce(total_test_correct_certain, op=dist.ReduceOp.SUM) dist.all_reduce(total_test_mazes_solved, op=dist.ReduceOp.SUM) dist.all_reduce(total_test_steps, op=dist.ReduceOp.SUM) dist.all_reduce(total_test_mazes, op=dist.ReduceOp.SUM) # Calculate and Plot final Test metrics on Rank 0 if is_main_process(rank) and total_test_mazes > 0: avg_test_loss = total_test_loss.item() / total_test_mazes.item() avg_test_acc_step = total_test_correct_certain.item() / total_test_steps.item() avg_test_acc_maze = total_test_mazes_solved.item() / total_test_mazes.item() test_losses.append(avg_test_loss) test_accuracies_most_certain.append(avg_test_acc_step) test_accuracies_most_certain_permaze.append(avg_test_acc_maze) print(f"Iter {bi} Test Metrics (Agg): Loss={avg_test_loss:.4f}, StepAcc={avg_test_acc_step:.4f}, MazeAcc={avg_test_acc_maze:.4f}\n") # --- Plotting --- figacc = plt.figure(figsize=(10, 10)) axacc_train = figacc.add_subplot(211) axacc_test = figacc.add_subplot(212) # Plot Avg Step Accuracy axacc_train.plot(iters, train_accuracies_most_certain, 'k-', alpha=0.7, label=f'Avg Step Acc ({train_accuracies_most_certain[-1]:.3f})') axacc_test.plot(iters, test_accuracies_most_certain, 'k-', alpha=0.7, label=f'Avg Step Acc ({test_accuracies_most_certain[-1]:.3f})') # Plot Full Maze Accuracy axacc_train.plot(iters, train_accuracies_most_certain_permaze, 'r-', alpha=0.6, label=f'Full Maze Acc ({train_accuracies_most_certain_permaze[-1]:.3f})') axacc_test.plot(iters, test_accuracies_most_certain_permaze, 'r-', alpha=0.6, label=f'Full Maze Acc ({test_accuracies_most_certain_permaze[-1]:.3f})') axacc_train.set_title('Train Accuracy (Aggregated)') axacc_test.set_title('Test Accuracy (Aggregated)') axacc_train.legend(loc='lower right') axacc_test.legend(loc='lower right') axacc_train.set_xlim([0, args.training_iterations]) axacc_test.set_xlim([0, args.training_iterations]) axacc_train.set_ylim([0, 1]) axacc_test.set_ylim([0, 1]) figacc.tight_layout() figacc.savefig(f'{args.log_dir}/accuracies.png', dpi=150) plt.close(figacc) # Loss Plot figloss = plt.figure(figsize=(10, 5)) axloss = figloss.add_subplot(111) axloss.plot(iters, train_losses, 'b-', linewidth=1, alpha=0.8, label=f'Train (Agg): {train_losses[-1]:.4f}') axloss.plot(iters, test_losses, 'r-', linewidth=1, alpha=0.8, label=f'Test (Agg): {test_losses[-1]:.4f}') axloss.legend(loc='upper right') axloss.set_xlabel("Iteration") axloss.set_ylabel("Loss") axloss.set_xlim([0, args.training_iterations]) axloss.set_ylim(bottom=0) figloss.tight_layout() figloss.savefig(f'{args.log_dir}/losses.png', dpi=150) plt.close(figloss) # --- End Plotting --- # --- Visualization (Rank 0, Conditional) --- if is_main_process(rank) and args.model in ['ctm', 'lstm']: # try: model_module = model.module if isinstance(model, DDP) else model # Use a consistent batch for viz if possible, or just next batch inputs_viz, targets_viz = next(iter(testloader)) inputs_viz = inputs_viz.to(device) targets_viz = targets_viz.to(device) longest_index = (targets_viz!=4).sum(-1).argmax() # 4 assumed padding pbar.set_description('Tracking (Rank 0): Viz Fwd Pass') predictions_viz_raw, _, _, _, post_activations_viz, attention_tracking_viz = model_module(inputs_viz, track=True) predictions_viz = predictions_viz_raw.reshape(predictions_viz_raw.size(0), -1, 5, predictions_viz_raw.size(-1)) att_shape = (model.module.kv_features.shape[2], model.module.kv_features.shape[3]) attention_tracking_viz = attention_tracking_viz.reshape( attention_tracking_viz.shape[0], attention_tracking_viz.shape[1], -1, att_shape[0], att_shape[1]) pbar.set_description('Tracking (Rank 0): Dynamics Plot') plot_neural_dynamics(post_activations_viz, 100, args.log_dir, axis_snap=True) pbar.set_description('Tracking (Rank 0): Maze GIF') if attention_tracking_viz is not None: make_maze_gif((inputs_viz[longest_index].detach().cpu().numpy()+1)/2, predictions_viz[longest_index].detach().cpu().numpy(), targets_viz[longest_index].detach().cpu().numpy(), attention_tracking_viz[:, longest_index], args.log_dir) # else: # print("Skipping maze GIF due to attention shape issue.") # except Exception as e_viz: # print(f"Rank 0 visualization failed: {e_viz}") # --- End Visualization --- gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() if world_size > 1: dist.barrier() model.train() # --- End Evaluation Block --- if hasattr(train_sampler, 'set_epoch'): # Check if sampler has set_epoch train_sampler.set_epoch(bi) current_lr = optimizer.param_groups[-1]['lr'] try: inputs, targets = next(iterator) except StopIteration: iterator = iter(trainloader) inputs, targets = next(iterator) inputs = inputs.to(device, non_blocking=True) targets = targets.to(device, non_blocking=True) # Defaults for logging loss = torch.tensor(0.0, device=device) # Need loss defined for logging scope accuracy_finegrained = 0.0 where_most_certain_val = -1.0 where_most_certain_std = 0.0 where_most_certain_min = -1 where_most_certain_max = -1 upto_where_mean = -1.0 upto_where_std = 0.0 upto_where_min = -1 upto_where_max = -1 with torch.autocast(device_type="cuda" if device.type == 'cuda' else "cpu", dtype=torch.float16, enabled=args.use_amp): if args.do_compile: torch.compiler.cudagraph_mark_step_begin() if args.model == 'ctm': predictions_raw, certainties, _ = model(inputs) predictions = predictions_raw.reshape(predictions_raw.size(0), -1, 5, predictions_raw.size(-1)) # B,S,C,T loss, where_most_certain, upto_where = maze_loss(predictions, certainties, targets, cirriculum_lookahead=args.cirriculum_lookahead, use_most_certain=True) with torch.no_grad(): # Calculate local accuracy for logging accuracy_finegrained = (predictions.argmax(2)[torch.arange(predictions.size(0), device=device), :, where_most_certain] == targets).float().mean().item() elif args.model == 'lstm': predictions_raw, certainties, _ = model(inputs) predictions = predictions_raw.reshape(predictions_raw.size(0), -1, 5, predictions_raw.size(-1)) # B,S,C,T loss, where_most_certain, upto_where = maze_loss(predictions, certainties, targets, cirriculum_lookahead=args.cirriculum_lookahead, use_most_certain=False) # where = -1 with torch.no_grad(): accuracy_finegrained = (predictions.argmax(2)[torch.arange(predictions.size(0), device=device), :, where_most_certain] == targets).float().mean().item() elif args.model == 'ff': predictions_raw = model(inputs) # B, S*C predictions = predictions_raw.reshape(predictions_raw.size(0), -1, 5) # B,S,C loss, where_most_certain, upto_where = maze_loss(predictions.unsqueeze(-1), None, targets, cirriculum_lookahead=args.cirriculum_lookahead, use_most_certain=False) # where = -1 with torch.no_grad(): accuracy_finegrained = (predictions.argmax(2) == targets).float().mean().item() # Extract stats from loss outputs if torch.is_tensor(where_most_certain): where_most_certain_val = where_most_certain.float().mean().item() where_most_certain_std = where_most_certain.float().std().item() where_most_certain_min = where_most_certain.min().item() where_most_certain_max = where_most_certain.max().item() elif isinstance(where_most_certain, int): where_most_certain_val = float(where_most_certain); where_most_certain_min = where_most_certain; where_most_certain_max = where_most_certain if isinstance(upto_where, (np.ndarray, list)) and len(upto_where) > 0: upto_where_mean = np.mean(upto_where); upto_where_std = np.std(upto_where); upto_where_min = np.min(upto_where); upto_where_max = np.max(upto_where) # Backprop / Step scaler.scale(loss).backward() if args.gradient_clipping!=-1: scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=args.gradient_clipping) scaler.step(optimizer) scaler.update() optimizer.zero_grad(set_to_none=True) scheduler.step() # --- Aggregation and Logging (Rank 0) --- loss_log = loss.detach() if world_size > 1: dist.all_reduce(loss_log, op=dist.ReduceOp.AVG) if is_main_process(rank): pbar_desc = f'Loss(avg)={loss_log.item():.3f} Acc(loc)={accuracy_finegrained:.3f} LR={current_lr:.6f}' if args.model in ['ctm', 'lstm'] or torch.is_tensor(where_most_certain): pbar_desc += f' Cert={where_most_certain_val:.2f}'#+-{where_most_certain_std:.2f}' # Removed std for brevity if isinstance(upto_where, (np.ndarray, list)) and len(upto_where) > 0: pbar_desc += f' Path={upto_where_mean:.1f}'#+-{upto_where_std:.1f}' pbar.set_description(f'{args.model.upper()} {pbar_desc}') # --- End Aggregation and Logging --- # --- Checkpointing (Rank 0) --- if (bi % args.save_every == 0 or bi == args.training_iterations - 1) and bi != start_iter and is_main_process(rank): pbar.set_description('Rank 0: Saving checkpoint...') save_path = f'{args.log_dir}/checkpoint.pt' model_state_to_save = model.module.state_dict() if isinstance(model, DDP) else model.state_dict() checkpoint_data = { 'model_state_dict': model_state_to_save, 'optimizer_state_dict': optimizer.state_dict(), 'scheduler_state_dict': scheduler.state_dict(), 'scaler_state_dict': scaler.state_dict(), 'iteration': bi, 'train_losses': train_losses, 'test_losses': test_losses, 'train_accuracies': train_accuracies, # Saving simplified scalar list 'test_accuracies': test_accuracies, # Saving simplified scalar list 'train_accuracies_most_certain': train_accuracies_most_certain, 'test_accuracies_most_certain': test_accuracies_most_certain, 'train_accuracies_most_certain_permaze': train_accuracies_most_certain_permaze, 'test_accuracies_most_certain_permaze': test_accuracies_most_certain_permaze, 'iters': iters, 'args': args, 'torch_rng_state': torch.get_rng_state(), 'numpy_rng_state': np.random.get_state(), 'random_rng_state': random.getstate(), } torch.save(checkpoint_data, save_path) # --- End Checkpointing --- if world_size > 1: dist.barrier() if is_main_process(rank): pbar.update(1) # --- End Training Loop --- if is_main_process(rank): pbar.close() cleanup_ddp()