tasks/sort/train.py

import argparse
import os
import random

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')   
from tqdm.auto import tqdm

from data.custom_datasets import SortDataset
from models.ctm_sort import ContinuousThoughtMachineSORT
from tasks.image_classification.plotting import plot_neural_dynamics, make_classification_gif
from utils.housekeeping import set_seed, zip_python_code
from utils.losses import sort_loss
from tasks.sort.utils import compute_ctc_accuracy, decode_predictions
from utils.schedulers import WarmupCosineAnnealingLR, WarmupMultiStepLR, warmup

import torchvision
torchvision.disable_beta_transforms_warning()

from autoclip.torch import QuantileClip

import warnings
warnings.filterwarnings("ignore", message="using precomputed metric; inverse_transform will be unavailable")


warnings.filterwarnings(
    "ignore",
    "Corrupt EXIF data",
    UserWarning,
    r"^PIL\.TiffImagePlugin$" # Using a regular expression to match the module.
)

warnings.filterwarnings(
    "ignore",
    "UserWarning: Metadata Warning",
    UserWarning,
    r"^PIL\.TiffImagePlugin$" # Using a regular expression to match the module.
)


warnings.filterwarnings(
    "ignore",
    "UserWarning: Truncated File Read",
    UserWarning,
    r"^PIL\.TiffImagePlugin$" # Using a regular expression to match the module.
)


def parse_args():
    parser = argparse.ArgumentParser()

    # Model Architecture
    parser.add_argument('--d_model', type=int, default=512, help='Dimension of the model.')
    parser.add_argument('--d_input', type=int, default=128, help='Dimension of the input.')
    parser.add_argument('--synapse_depth', type=int, default=4, help='Depth of U-NET model for synapse. 1=linear, no unet.')
    parser.add_argument('--heads', type=int, default=4, help='Number of attention heads.')
    parser.add_argument('--n_synch_out', type=int, default=32, help='Number of neurons to use for output synch.')
    parser.add_argument('--n_synch_action', type=int, default=32, help='Number of neurons to use for observation/action synch.')
    parser.add_argument('--neuron_select_type', type=str, default='random-pairing', help='Protocol for selecting neuron subset.')
    parser.add_argument('--n_random_pairing_self', type=int, default=0, help='Number of neurons paired self-to-self for synch.')
    
    parser.add_argument('--iterations', type=int, default=50, help='Number of internal ticks.')
    parser.add_argument('--memory_length', type=int, default=25, help='Length of the pre-activation history for NLMS.')
    parser.add_argument('--deep_memory', action=argparse.BooleanOptionalAction, default=True,
                        help='Use deep memory.')
    parser.add_argument('--memory_hidden_dims', type=int, default=4, help='Hidden dimensions of the memory if using deep memory.')
    parser.add_argument('--dropout', type=float, default=0.0, help='Dropout rate.')
    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.')
    parser.add_argument('--do_normalisation', action=argparse.BooleanOptionalAction, default=False,
                        help='Apply normalization in NLMs.')
    parser.add_argument('--positional_embedding_type', type=str, default='none',
                        help='Type of positional embedding.', choices=['none', 
                                                                       'learnable-fourier', 
                                                                       'multi-learnable-fourier',
                                                                       'custom-rotational'])

    # Training
    parser.add_argument('--batch_size', type=int, default=32, help='Batch size for training.')
    parser.add_argument('--batch_size_test', type=int, default=32, help='Batch size for testing.')
    parser.add_argument('--lr', type=float, default=1e-3, 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('--gradient_clipping', type=float, default=-1, help='Gradient quantile clipping value (-1 to disable).')
    parser.add_argument('--use_amp', action=argparse.BooleanOptionalAction, default=False, help='AMP autocast.')
    parser.add_argument('--do_compile', action=argparse.BooleanOptionalAction, default=False, help='Try to compile the synapses, backbone, and nlms.')


    # Logging and Saving
    parser.add_argument('--log_dir', type=str, default='logs/scratch',
                        help='Directory for logging.')
    parser.add_argument('--N_to_sort', type=int, default=30, help='N numbers to sort.')
    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?')

    # 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')

    # Device
    parser.add_argument('--device', type=int, nargs='+', default=[-1],
                        help='List of GPU(s) to use. Set to -1 to use CPU.')

    args = parser.parse_args()
    return args




if __name__=='__main__':

    # Hosuekeeping
    args = parse_args()
    # Change the following for sorting
    args.backbone_type = 'none'
    
    set_seed(args.seed, False)
    if not os.path.exists(args.log_dir): os.makedirs(args.log_dir)
    
    

    

    # Data
    train_data = SortDataset(args.N_to_sort)
    test_data = SortDataset(args.N_to_sort)
    trainloader = torch.utils.data.DataLoader(train_data, batch_size=args.batch_size, shuffle=True, num_workers=1)
    testloader = torch.utils.data.DataLoader(test_data, batch_size=args.batch_size_test, shuffle=True, num_workers=1, drop_last=False)
    

    prediction_reshaper = [-1]  # Problem specific
    args.out_dims = args.N_to_sort + 1

    # For total reproducibility
    # Python 3.x
    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)  

    # Configure device string (support MPS on macOS)
    if args.device[0] != -1:
        device = f'cuda:{args.device[0]}'
    elif torch.backends.mps.is_available():
        device = 'mps'
    else:
        device = 'cpu'
    print(f'Running model {args.model} on {device}')

    # Build model
    model = ContinuousThoughtMachineSORT(
        iterations=args.iterations,
        d_model=args.d_model,
        d_input=args.out_dims-1,  
        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='none',
        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)

    
    model.train()

    # For lazy modules so that we can get param count
    pseudo_inputs = train_data.__getitem__(0)[0].unsqueeze(0).to(device)
    model(pseudo_inputs)  

    print(f'Total params: {sum(p.numel() for p in model.parameters())}')
    
    

    # 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):
        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 (I like custom)
    # Using batched estimates
    start_iter = 0  # For reloading, keep track of this (pretty tqdm stuff needs it)
    train_losses = []  
    test_losses = []
    train_accuracies = []  # This will be per internal tick, not so simple
    test_accuracies = []
    train_accuracies_full_list = []  # This will be selected according to what is returned by loss function
    test_accuracies_full_list = []
    iters = []

    # Now that everything is initliased, reload if desired
    scaler = torch.amp.GradScaler("cuda" if "cuda" in device else "cpu", enabled=args.use_amp)
    if args.reload:
        if os.path.isfile(f'{args.log_dir}/checkpoint.pt'):
            print(f'Reloading from: {args.log_dir}/checkpoint.pt')
            checkpoint = torch.load(f'{args.log_dir}/checkpoint.pt', map_location=device, weights_only=False)
            model.load_state_dict(checkpoint['model_state_dict'], strict=True)
            if not args.reload_model_only:
                print('Reloading optimizer etc.')
                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']
                train_losses = checkpoint['train_losses']
                train_accuracies_full_list = checkpoint['train_accuracies_full_list']
                train_accuracies = checkpoint['train_accuracies']
                test_losses = checkpoint['test_losses']
                test_accuracies_full_list = checkpoint['test_accuracies_full_list']
                test_accuracies = checkpoint['test_accuracies']
                iters = checkpoint['iters']
            else:
                print('Only relading model!')
            if 'torch_rng_state' in checkpoint:
                # Reset seeds, otherwise mid-way training can be obscure (particularly for imagenet)
                torch.set_rng_state(checkpoint['torch_rng_state'].cpu().byte())
                np.random.set_state(checkpoint['numpy_rng_state'])
                random.setstate(checkpoint['random_rng_state'])

            del checkpoint
            import gc
            gc.collect()
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
    
    if args.do_compile:
        print('Compiling...')
        model.synapses = torch.compile(model.synapses, mode='reduce-overhead', fullgraph=True)
        model.backbone = torch.compile(model.backbone, mode='reduce-overhead', fullgraph=True)
    
    # Training
    iterator = iter(trainloader)  # Not training in epochs, but rather iterations. Need to reset this from time to time
    with tqdm(total=args.training_iterations, initial=start_iter, leave=False, position=0, dynamic_ncols=True) as pbar:
        for bi in range(start_iter, args.training_iterations):
            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)
            targets = targets.to(device)
            with torch.autocast(device_type="cuda" if "cuda" in device else "cpu", dtype=torch.float16, enabled=args.use_amp):
                if args.do_compile:
                    torch.compiler.cudagraph_mark_step_begin()
                predictions, certainties, synchronisation = model(inputs)
                loss = sort_loss(predictions, targets)
            
                        
            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()

            accuracy = compute_ctc_accuracy(predictions, targets, predictions.shape[1]-1)
            pbar.set_description(f'Sorting {args.N_to_sort} real numbers. Loss={loss.item():0.3f}. Accuracy={accuracy:0.3f}. LR={current_lr:0.6f}')


            # Metrics tracking and plotting
            if bi%args.track_every==0:# and bi != 0:
                model.eval()
                with torch.inference_mode():
                    

                    inputs, targets = next(iter(testloader))
                    inputs = inputs.to(device)
                    targets = targets.to(device)
                    pbar.set_description('Tracking: Processing test data')
                    predictions, certainties, synchronisation, pre_activations, post_activations, _ = model(inputs, track=True)
                    pbar.set_description('Tracking: Neural dynamics')
                    plot_neural_dynamics(post_activations, 100, args.log_dir)

                    imgi = 0


                    
                    ##################################### TRAIN METRICS
                    all_predictions = []
                    all_targets = []
                    all_losses = []
                    
                    iters.append(bi)
                    pbar.set_description('Tracking: Computing loss and accuracy for curves')
                    with torch.inference_mode():
                        loader = torch.utils.data.DataLoader(train_data, batch_size=args.batch_size_test, shuffle=True, num_workers=1)
                        with tqdm(total=len(loader), initial=0, leave=False, position=1, dynamic_ncols=True) as pbar_inner:
                        
                            for inferi, (inputs, targets) in enumerate(loader):
                                
                                inputs = inputs.to(device)
                                targets = targets.to(device)
                                these_predictions, certainties, synchronisation = model(inputs)

                                loss = sort_loss(these_predictions, targets)
                                all_losses.append(loss.item())

                                all_targets.append(targets.detach().cpu().numpy())

                                decoded = [d[:targets.shape[1]] for d in decode_predictions(these_predictions, predictions.shape[1]-1)]
                                decoded = torch.stack([torch.concatenate((d, torch.zeros(targets.shape[1] - len(d), device=targets.device)+targets.shape[1])) if len(d) < targets.shape[1] else d for d in decoded], 0)

                                all_predictions.append(decoded.detach().cpu().numpy())
                                
                                if args.n_test_batches!=-1 and inferi%args.n_test_batches==0 and inferi!=0 : break
                                pbar_inner.set_description('Computing metrics for train')
                                pbar_inner.update(1)

                        all_predictions = np.concatenate(all_predictions)
                        all_targets = np.concatenate(all_targets)


                        train_accuracies.append((all_predictions==all_targets).mean())
                        train_accuracies_full_list.append((all_predictions==all_targets).all(-1).mean())
                        train_losses.append(np.mean(all_losses))

                        ##################################### TEST METRICS
                        all_predictions = []
                        all_targets = []
                        all_losses = []
                        loader = torch.utils.data.DataLoader(test_data, batch_size=args.batch_size_test, shuffle=True, num_workers=1)
                        with tqdm(total=len(loader), initial=0, leave=False, position=1, dynamic_ncols=True) as pbar_inner:
                            for inferi, (inputs, targets) in enumerate(loader):
                                
                                inputs = inputs.to(device)
                                targets = targets.to(device)
                                these_predictions, certainties, synchronisation = model(inputs)

                                loss = sort_loss(these_predictions, targets)
                                all_losses.append(loss.item())

                                all_targets.append(targets.detach().cpu().numpy())

                                decoded = [d[:targets.shape[1]] for d in decode_predictions(these_predictions, predictions.shape[1]-1)]
                                decoded = torch.stack([torch.concatenate((d, torch.zeros(targets.shape[1] - len(d), device=targets.device)+targets.shape[1])) if len(d) < targets.shape[1] else d for d in decoded], 0)

                                all_predictions.append(decoded.detach().cpu().numpy())
                                
                                if args.n_test_batches!=-1 and inferi%args.n_test_batches==0 and inferi!=0 : break
                                pbar_inner.set_description('Computing metrics for train')
                                pbar_inner.update(1)

                        all_predictions = np.concatenate(all_predictions)
                        all_targets = np.concatenate(all_targets)


                        test_accuracies.append((all_predictions==all_targets).mean())
                        test_accuracies_full_list.append((all_predictions==all_targets).all(-1).mean())
                        test_losses.append(np.mean(all_losses))
                            

                        figacc = plt.figure(figsize=(10, 10))
                        axacc_train = figacc.add_subplot(211)
                        axacc_test = figacc.add_subplot(212)
                        cm = sns.color_palette("viridis", as_cmap=True)
                        axacc_train.plot(iters, train_accuracies, 'b-', alpha=0.7, label='Find grained')   
                        axacc_train.plot(iters, train_accuracies_full_list, 'k--', alpha=0.7, label='Full list')   
                        axacc_test.plot(iters, test_accuracies, 'b-', alpha=0.7, label='Fine grained')        
                        axacc_test.plot(iters, test_accuracies_full_list, 'k--', alpha=0.7, label='Full list')        
                        axacc_train.set_title('Train')
                        axacc_test.set_title('Test')
                        axacc_train.legend(loc='lower right')
                        axacc_train.set_xlim([0, args.training_iterations])
                        axacc_test.set_xlim([0, args.training_iterations])
                        
                        figacc.tight_layout()
                        figacc.savefig(f'{args.log_dir}/accuracies.png', dpi=150)
                        plt.close(figacc)

                        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: {train_losses[-1]}')
                        axloss.plot(iters, test_losses, 'r-', linewidth=1, alpha=0.8, label=f'Test: {test_losses[-1]}')
                        axloss.legend(loc='upper right')

                        axloss.set_xlim([0, args.training_iterations])
                        figloss.tight_layout()
                        figloss.savefig(f'{args.log_dir}/losses.png', dpi=150)
                        plt.close(figloss)

                model.train()
                            



            # Save model
            if (bi%args.save_every==0 or bi==args.training_iterations-1) and bi != start_iter:
                torch.save(
                    {
                    'model_state_dict':model.state_dict(),
                    'optimizer_state_dict':optimizer.state_dict(),
                    'scheduler_state_dict':scheduler.state_dict(),
                    'scaler_state_dict':scaler.state_dict(),
                    'iteration':bi,
                    'train_accuracies_full_list':train_accuracies_full_list,
                    'train_accuracies':train_accuracies,
                    'test_accuracies_full_list':test_accuracies_full_list,
                    'test_accuracies':test_accuracies,
                    'train_losses':train_losses,
                    'test_losses':test_losses,
                    'iters':iters,
                    'args':args,
                    'torch_rng_state': torch.get_rng_state(),
                    'numpy_rng_state': np.random.get_state(),
                    'random_rng_state': random.getstate(),
                    } , f'{args.log_dir}/checkpoint.pt')
            
            pbar.update(1)