source/official-code/micro/train_utils.py

import torch
import torch.optim as optim
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader

import json
import os
import numpy as np
import math

from generate import generate_data
from models.hybrid import HybridModel
from models.transformer import generate_mask



def get_ident_name(args):
    dashed_task_name = "-".join(args.task_name.split("_"))
    return "run_%s_%s_%s_w%d_d%d_nh%d_sd%d_nn%d_nv%d" % (dashed_task_name, args.layer1, args.layer2, args.window, args.embed_dim, \
                                           args.num_heads, args.state_dim, args.num_numbers, args.num_vocab)

def get_task_dir_name(args):
    return args.task_name + "/" + args.data_name + "/" + get_ident_name(args)

def make_dir(args):
    if args.ood_eval:
        if 'results_ood' not in os.listdir('.'):
            os.mkdir('results_ood')
        base_path = 'results_ood'
    else:
        if 'results' not in os.listdir('.'):
            os.mkdir('results')
        base_path = 'results'
    
    if args.task_name not in os.listdir(base_path + ''):
        os.mkdir(base_path + '/' + args.task_name)

    if args.data_name not in os.listdir(base_path + '/' + args.task_name):
        os.mkdir(base_path + '/' + args.task_name + "/" + args.data_name)
    
    if get_ident_name(args) not in os.listdir(base_path + '/' + args.task_name + "/" + args.data_name):
        os.mkdir(base_path + '/' + args.task_name + "/" + args.data_name + "/" + get_ident_name(args))



class CustomDataset(Dataset):
    def __init__(self, data_in, data_out):
        self.data_in = data_in
        self.data_out = data_out

    def __len__(self):
        return len(self.data_in)

    def __getitem__(self, idx):
        sample_in = self.data_in[idx]
        sample_out = self.data_out[idx]

        return sample_in, sample_out
    


def get_scheduler(args, optimizer):
    total_steps = args.batches_per_epoch * args.num_epochs
    warmup_steps = args.batches_per_epoch // 10 # Make the warmup be 10% of an epoch
    
    def lr_lambda(step):
        if step < warmup_steps:
            return step / warmup_steps
        progress = (step - warmup_steps) / (total_steps - warmup_steps)
        return 0.5 * (1 + math.cos(math.pi * progress))

    return optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)



def train_epoch(model, optimizer, lr_scheduler, criterion, mask, train_loader, args, device="cuda"):
    model.train()

    avg_loss = 0
    avg_acc = 0
    
    # for itr in range(args.batches_per_epoch):
    #     print("Batch:", itr)
    #     input_seqs, target_seqs = generate_data(args)
    
    for itr, (input_seqs, target_seqs) in enumerate(train_loader):
        
        input_seqs = input_seqs.to(device)
        target_seqs = target_seqs.type(torch.LongTensor).to(device)

        outputs = model(input_seqs, mask)  # (batch, seq-1, vocab)

        # Masked loss, ignore positions which have null tokens (== vocab_size-1)
        loss_mask = (target_seqs != args.vocab_size-1)
        loss = loss_mask.reshape(-1) * criterion(outputs.view(-1, args.vocab_size), target_seqs.reshape(-1))
        loss = loss.sum() / loss_mask.sum()

        acc = torch.sum(loss_mask & ((torch.argmax(outputs, dim=-1) - target_seqs) == 0)).item()
        acc /= loss_mask.sum()

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        lr_scheduler.step()

        avg_loss += loss.item()
        avg_acc += acc

    avg_loss /= args.batches_per_epoch
    avg_acc /= args.batches_per_epoch
    
    return avg_loss, avg_acc


def eval_epoch(model, criterion, mask, eval_loader, args, device="cuda"):
    model.eval()

    avg_loss = 0
    avg_acc = 0
    
    for itr, (input_seqs, target_seqs) in enumerate(eval_loader):
        
        input_seqs = input_seqs.to(device)
        target_seqs = target_seqs.type(torch.LongTensor).to(device)

        outputs = model(input_seqs, mask)  # (batch, seq-1, vocab)

        # Masked loss, ignore positions which have null tokens (== vocab_size-1)
        loss_mask = (target_seqs != args.vocab_size-1)
        loss = loss_mask.reshape(-1) * criterion(outputs.view(-1, args.vocab_size), target_seqs.reshape(-1))
        loss = loss.sum() / loss_mask.sum()

        acc = torch.sum(loss_mask & ((torch.argmax(outputs, dim=-1) - target_seqs) == 0)).item()
        acc /= loss_mask.sum()

        avg_loss += loss.item()
        avg_acc += acc

    avg_loss /= args.batches_per_epoch
    avg_acc /= args.batches_per_epoch
    
    return avg_loss, avg_acc



def choose_lr(args):
    # If found already, use what is cached
    ident_name = get_ident_name(args)
    with open('models/lrs.json') as f:
        d = json.load(f)
    if (not args.lr_no_save) and ident_name in d.keys():
        return d[ident_name]

    train_dataset = CustomDataset(*generate_data(args, all_at_once=True))
    train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True)
    
    assert args.window <= args.sequence_len
    
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
    print("Starting to find lr for results/" + args.data_name + "/" + args.run_name)

    lrs = np.geomspace(args.lr_low, args.lr_high, args.lr_num)
    criterion = nn.CrossEntropyLoss(reduction='none')
    all_losses = np.zeros_like(lrs)

    for itr, lr in enumerate(lrs):
        curr_losses = []
        print("- Current learning rate: %.4f" % lr)
        for run in range(3):
            model = HybridModel(args).to(device)
        
            optimizer = optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-2)
            lr_scheduler = get_scheduler(args, optimizer)
        
            mask = generate_mask(args.sequence_len, window=args.window).to(device)
        
            # Training loop
            for epoch in range(args.lr_epochs):
                loss, acc = train_epoch(model, optimizer, lr_scheduler, criterion, mask, train_loader, args, device)
        
            print(f"- - Loss: {loss:.4f}, Acc: {acc:.4f}")
            
            curr_losses.append(loss)
            
        all_losses[itr] = np.median(np.array(curr_losses))

    all_losses = np.nan_to_num(all_losses, nan=100) # Nan needs to be a big number
    lr = lrs[np.argmin(all_losses)]
    print("Loss array:", all_losses)
    print("Determined a lr of", lr)

    if not args.lr_no_save:
        # Saving the lr
        with open('models/lrs.json') as f:
            d = json.load(f)
        d[ident_name] = lr
        with open('models/lrs.json', 'w') as f:
            json.dump(d, f, indent=4)

    return lr


    
def train(args):
    assert args.window <= args.sequence_len

    if args.save:
        dir_name = get_task_dir_name(args)
        if not args.force_learn:
            if args.ood_eval:
                if os.path.isdir("results_ood/" + dir_name) and "run%d.pt" % args.run_number in os.listdir("results_ood/" + dir_name):
                    return
            else:
                if os.path.isdir("results/" + dir_name) and "run%d.pt" % args.run_number in os.listdir("results/" + dir_name):
                    return

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    train_dataset = CustomDataset(*generate_data(args, all_at_once=True))
    train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True)
    
    print("Starting work for results/" + args.data_name + "/" + args.run_name)

    model = HybridModel(args).to(device)
    criterion = nn.CrossEntropyLoss(reduction='none')

    optimizer = optim.AdamW(model.parameters(), lr=args.lr, weight_decay=1e-2)
    lr_scheduler = get_scheduler(args, optimizer)

    mask = generate_mask(args.sequence_len, window=args.window).to(device)

    # Training loop
    losses = []
    accs = []
        
    for epoch in range(args.num_epochs):
        loss, acc = train_epoch(model, optimizer, lr_scheduler, criterion, mask, train_loader, args, device)

        losses.append(loss)
        accs.append(acc)

    print(f"Loss: {losses[-1]:.4f}, Acc: {accs[-1]:.4f}")
    
    losses = torch.Tensor(losses)
    accs = torch.Tensor(accs)

    # Evaluation
    if args.ood_eval:
        args.num_vocab = args.eval_num_vocab
        args.num_numbers = args.eval_num_numbers
        args.p = args.eval_p
        args.num_bits = args.eval_num_bits

    eval_dataset = CustomDataset(*generate_data(args, all_at_once=True))
    eval_loader = DataLoader(eval_dataset, batch_size=args.batch_size, shuffle=True)
    
    eval_loss, eval_acc = eval_epoch(model, criterion, mask, eval_loader, args, device)
    print(f"Eval - Loss: {eval_loss:.4f}, Acc: {eval_acc:.4f}")
    
    # Save the model
    if args.save and not args.train_only:
        make_dir(args)
        dir_name = get_task_dir_name(args)
        if args.ood_eval:
            run_filename = "results_ood/" + dir_name + "/run%d.pt" % args.run_number
        else:
            run_filename = "results/" + dir_name + "/run%d.pt" % args.run_number
        
        # Count of the parameters
        params = sum(p.numel() for p in model.parameters() if p.requires_grad)

        torch.save({"args": args, "losses": losses, "accs": accs, "param_count": params, "eval_loss": eval_loss, "eval_acc": eval_acc}, run_filename)
        print("Saved and finished for %s" % run_filename)

    return model