TabPFN/train.py

import os
import itertools
import argparse
import time
import datetime
import yaml
from contextlib import nullcontext


import torch
from torch import nn

import utils
from transformer import TransformerModel
from utils import get_cosine_schedule_with_warmup, get_openai_lr, StoreDictKeyPair, get_weighted_single_eval_pos_sampler, get_uniform_single_eval_pos_sampler
import priors
import encoders
import positional_encodings
from utils import init_dist
from torch.cuda.amp import autocast

class Losses():
    gaussian = nn.GaussianNLLLoss(full=True, reduction='none')
    mse = nn.MSELoss(reduction='none')
    ce = lambda weight : nn.CrossEntropyLoss(reduction='none', weight=weight)
    bce = nn.BCEWithLogitsLoss(reduction='none')


def train(priordataloader_class, criterion, encoder_generator, emsize=200, nhid=200, nlayers=6, nhead=2, dropout=0.2,
          epochs=10, steps_per_epoch=100, batch_size=200, bptt=10, lr=None, weight_decay=0.0, warmup_epochs=10, input_normalization=False,
          y_encoder_generator=None, pos_encoder_generator=None, decoder=None, extra_prior_kwargs_dict={}, scheduler=get_cosine_schedule_with_warmup,
          load_weights_from_this_state_dict=None, validation_period=10, single_eval_pos_gen=None, bptt_extra_samples=None, gpu_device='cuda:0',
          aggregate_k_gradients=1, verbose=True, style_encoder_generator=None, check_is_compatible=True, epoch_callback=None,
          initializer=None, initialize_with_model=None, train_mixed_precision=False, total_available_time_in_s=None, normalize_labels=True, **model_extra_args
          ):
    assert (epochs is None) != (total_available_time_in_s is None)
    start_of_training = time.time()
    device = gpu_device if torch.cuda.is_available() else 'cpu:0'
    print(f'Using {device} device')
    using_dist, rank, device = init_dist(device)
    bptt_sampler = (lambda : single_eval_pos_gen() + bptt_extra_samples if callable(single_eval_pos_gen) else single_eval_pos_gen + bptt_extra_samples) if bptt_extra_samples is not None else bptt
    dl = priordataloader_class(num_steps=steps_per_epoch, batch_size=batch_size, seq_len=bptt_sampler, seq_len_maximum=bptt+(bptt_extra_samples if bptt_extra_samples else 0), device=device, **extra_prior_kwargs_dict)
    if dl.fuse_x_y:
        raise Exception("Illegal parameter")

    encoder = encoder_generator(dl.num_features+1 if dl.fuse_x_y else dl.num_features,emsize)
    style_def = next(iter(dl))[0][0] # This is (style, x, y), target with x and y with batch size

    style_encoder = style_encoder_generator(hyperparameter_definitions=style_def[0], em_size=emsize) if (style_def is not None) else None
    n_out = dl.num_outputs
    if isinstance(criterion, nn.GaussianNLLLoss):
        n_out *= 2
    elif isinstance(criterion, nn.CrossEntropyLoss):
        n_out *= criterion.weight.shape[0]
    model = TransformerModel(encoder, n_out, emsize, nhead, nhid, nlayers, dropout, style_encoder=style_encoder,
                             y_encoder=y_encoder_generator(dl.num_outputs, emsize), input_normalization=input_normalization,
                             pos_encoder=(pos_encoder_generator or positional_encodings.NoPositionalEncoding)(emsize, bptt*2),
                             decoder=decoder, init_method=initializer, **model_extra_args
                             )
    model.criterion = criterion
    if load_weights_from_this_state_dict is not None:
        model.load_state_dict(load_weights_from_this_state_dict)
    if initialize_with_model is not None:
        model.init_from_small_model(initialize_with_model)

    print(f"Using a Transformer with {sum(p.numel() for p in model.parameters())/1000/1000:.{2}f} M parameters")

    try:
        for (k, v), (k2, v2) in zip(model.state_dict().items(), initialize_with_model.state_dict().items()):
            print(k, ((v - v2) / v).abs().mean(), v.shape)
    except Exception:
        pass

    model.to(device)
    if using_dist:
        print("Distributed training")
        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[rank], output_device=rank, broadcast_buffers=False)


    # learning rate
    if lr is None:
        lr = get_openai_lr(model)
        print(f"Using OpenAI max lr of {lr}.")
    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=weight_decay)
    scheduler = scheduler(optimizer, warmup_epochs, epochs if epochs is not None else 100) # when training for fixed time lr schedule takes 100 steps

    def train_step():
        model.train()  # Turn on the train mode
        total_loss = 0.
        total_positional_losses = 0.
        total_positional_losses_recorded = 0
        before_get_batch = time.time()
        assert len(dl) % aggregate_k_gradients == 0, 'Please set the number of steps per epoch s.t. `aggregate_k_gradients` divides it.'
        valid_batch_steps = 0.0
        for batch, (data, targets) in enumerate(dl):
            if using_dist and not (batch % aggregate_k_gradients == aggregate_k_gradients - 1):
                cm = model.no_sync()
                #print(f'p={rank}, no_sync', force=True)
            else:
                cm = nullcontext()
                #print(f'p={rank}, sync', force=True)
            with cm:
                time_to_get_batch = time.time() - before_get_batch
                before_forward = time.time()
                if bptt_extra_samples is None:
                    single_eval_pos = single_eval_pos_gen() if callable(single_eval_pos_gen) else single_eval_pos_gen
                else:
                    single_eval_pos = targets.shape[0] - bptt_extra_samples

                is_compatible = torch.ones((targets.shape[1])).bool()
                if check_is_compatible or normalize_labels:
                    for b in range(targets.shape[1]):
                        targets_in_train = torch.unique(targets[:single_eval_pos, b], sorted=True)
                        targets_in_eval = torch.unique(targets[single_eval_pos:, b], sorted=True)

                        if check_is_compatible:
                            is_compatible[b] = len(targets_in_train) == len(targets_in_eval) and (targets_in_train == targets_in_eval).all()
                            is_compatible[b] = is_compatible[b] and len(targets_in_train) > 1

                        # Set targets to range starting from 0 (e.g. targets 0, 2, 5, 2 will be converted to 0, 1, 2, 1)
                        if normalize_labels:
                            targets[:, b] = (targets[:, b] > torch.unique(targets[:, b]).unsqueeze(1)).sum(axis=0).unsqueeze(0)
                valid_batch_steps += is_compatible.float().mean()
                is_compatible = is_compatible.to(device)
                #if using_dist and check_is_compatible:
                #    print('step share before reduce',curr_step_share, force=True)
                #    curr_step_share = curr_step_share.to(device)
                #    torch.distributed.all_reduce_multigpu([curr_step_share], op=torch.distributed.ReduceOp.SUM)
                #    curr_step_share = curr_step_share.cpu() / torch.distributed.get_world_size()
                #    print('step share after reduce',curr_step_share, torch.distributed.get_world_size(), force=True)

                # If style is set to None, it should not be transferred to device
                output = model(tuple(e.to(device) if torch.is_tensor(e) else e for e in data) if isinstance(data, tuple) else data.to(device)
                               , single_eval_pos=single_eval_pos)

                forward_time = time.time() - before_forward

                #output, targets = output[:, is_compatible], targets[:, is_compatible]

                if single_eval_pos is not None:
                    targets = targets[single_eval_pos:]
                if isinstance(criterion, nn.GaussianNLLLoss):
                    assert output.shape[-1] == 2, \
                        'need to write a little bit of code to handle multiple regression targets at once'

                    mean_pred = output[..., 0]
                    var_pred = output[..., 1].abs()
                    losses = criterion(mean_pred.flatten(), targets.to(device).flatten(), var=var_pred.flatten())
                elif isinstance(criterion, (nn.MSELoss, nn.BCEWithLogitsLoss)):
                    losses = criterion(output.flatten(), targets.to(device).flatten())
                elif isinstance(criterion, (nn.CrossEntropyLoss)):
                    #print(n_out, targets.min(), targets.max(), force=True)
                    losses = criterion(output.reshape(-1, n_out), targets.to(device).long().flatten())
                else:
                    losses = criterion(output.reshape(-1, n_out), targets.to(device).flatten())
                losses = losses.view(*output.shape[0:2])
                loss = losses.mean(0) @ is_compatible.float() / losses.shape[1]
                #loss = torch_nanmean(losses, axis=[0, 1]) * is_compatible.float().mean()
                # not sure whether we can go without the nan checks.

                loss.backward()

                if ((batch % aggregate_k_gradients == aggregate_k_gradients - 1) and (not check_is_compatible or using_dist))\
                        or (valid_batch_steps >= aggregate_k_gradients and (check_is_compatible and not using_dist)):
                    with torch.no_grad():
                        for p in model.parameters():
                            if p.grad is not None:
                                p.grad.div_(valid_batch_steps)
                    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.)
                    try:
                        optimizer.step()
                    except:
                        print("Invalid optimization step encountered")
                    optimizer.zero_grad()
                    valid_batch_steps = 0.0

                step_time = time.time() - before_forward

                if not torch.isnan(loss):
                    total_loss += loss.item()
                    total_positional_losses += losses.mean(1).cpu().detach() if single_eval_pos is None else \
                        nn.functional.one_hot(torch.tensor(single_eval_pos), bptt)*loss.cpu().detach()

                    total_positional_losses_recorded += torch.ones(bptt) if single_eval_pos is None else \
                        nn.functional.one_hot(torch.tensor(single_eval_pos), bptt)

            before_get_batch = time.time()
        return total_loss / steps_per_epoch, (
                    total_positional_losses / total_positional_losses_recorded).tolist(), time_to_get_batch, forward_time, step_time

    best_val_loss = float("inf")
    best_model = None
    total_loss = float('inf')
    total_positional_losses = float('inf')
    try:
        for epoch in (range(1, epochs + 1) if epochs is not None else itertools.count(1)):

            epoch_start_time = time.time()
            if train_mixed_precision:
                with autocast():
                    total_loss, total_positional_losses, time_to_get_batch, forward_time, step_time = train_step()
            else:
                total_loss, total_positional_losses, time_to_get_batch, forward_time, step_time = train_step()
            if hasattr(dl, 'validate') and epoch % validation_period == 0:
                with torch.no_grad():
                    val_score = dl.validate(model)
            else:
                val_score = None

            if verbose:
                print('-' * 89)
                print(
                    f'| end of epoch {epoch:3d} | time: {(time.time() - epoch_start_time):5.2f}s | mean loss {total_loss:5.2f} | '
                    f"pos losses {','.join([f'{l:5.2f}' for l in total_positional_losses])}, lr {scheduler.get_last_lr()[0]}"
                    f' data time {time_to_get_batch:5.2f} step time {step_time:5.2f}'
                    f' forward time {forward_time:5.2f}' + (f'val score {val_score}' if val_score is not None else ''))
                print('-' * 89)

            # stepping with wallclock time based scheduler
            current_time = time.time()
            if epoch_callback is not None and rank == 0:
                epoch_callback(model, epoch / epochs if total_available_time_in_s is None else  # noqa
                (current_time - start_of_training) / total_available_time_in_s  # noqa
                               )
            if epochs is None and (current_time - start_of_training) > total_available_time_in_s:  # noqa
                break
            if epochs is None:
                scheduler.step((current_time - epoch_start_time) / total_available_time_in_s * 100)
            else:
                scheduler.step()
    except KeyboardInterrupt:
        pass

    return total_loss, total_positional_losses, model.to('cpu'), dl

def _parse_args(config_parser, parser):
    # Do we have a config file to parse?
    args_config, remaining = config_parser.parse_known_args()
    if args_config.config:
        with open(args_config.config, 'r') as f:
            cfg = yaml.safe_load(f)
            parser.set_defaults(**cfg)

    # The main arg parser parses the rest of the args, the usual
    # defaults will have been overridden if config file specified.
    args = parser.parse_args(remaining)

    # Cache the args as a text string to save them in the output dir later
    args_text = yaml.safe_dump(args.__dict__, default_flow_style=False)
    return args, args_text


if __name__ == '__main__':
    config_parser = argparse.ArgumentParser(description='Only used as a first parser for the config file path.')
    config_parser.add_argument('--config')
    parser = argparse.ArgumentParser()
    parser.add_argument('prior')
    parser.add_argument('--loss_function', default='barnll')
    # Optional Arg's for `--loss_function barnll`
    parser.add_argument('--min_y', type=float, help='barnll can only model y in strict ranges, this is the minimum y can take.')
    parser.add_argument('--max_y', type=float, help='barnll can only model y in strict ranges, this is the maximum y can take.')
    parser.add_argument('--num_buckets', default=100, type=int)
    #parser.add_argument('--num_features', default=None, type=int, help='Specify depending on the prior.')
    parser.add_argument("--extra_prior_kwargs_dict", default={'fuse_x_y': False}, dest="extra_prior_kwargs_dict", action=StoreDictKeyPair, nargs="+", metavar="KEY=VAL", help='Specify depending on the prior.')
    parser.add_argument('--encoder', default='linear', type=str, help='Specify depending on the prior.')
    parser.add_argument('--y_encoder', default='linear', type=str, help='Specify depending on the prior. You should specify this if you do not fuse x and y.')
    parser.add_argument('--pos_encoder', default='sinus', type=str, help='Specify depending on the prior.')
    parser.add_argument('--bptt', default=10, type=int)
    parser.add_argument('--epochs', default=200, type=int)
    parser.add_argument('--warmup_epochs', default=50, type=int)
    parser.add_argument('--validation_period', default=10, type=int)
    parser.add_argument('--permutation_invariant_max_eval_pos', default=None, type=int, help='Set this to an int to ')
    parser.add_argument('--permutation_invariant_sampling', default='weighted', help="Only relevant if --permutation_invariant_max_eval_pos is set.")

    # these can likely be mostly left at defaults
    parser.add_argument('--emsize', default=512, type=int) # sometimes even larger is better e.g. 1024
    parser.add_argument('--nlayers', default=6, type=int)
    parser.add_argument('--nhid', default=None, type=int) # 2*emsize is the default
    parser.add_argument('--nhead', default=4, type=int) # nhead = emsize / 64 in the original paper
    parser.add_argument('--dropout', default=.0, type=float)
    parser.add_argument('--steps_per_epoch', default=10, type=int)
    parser.add_argument('--batch_size', default=1000, type=int)
    parser.add_argument('--lr', '--learning_rate', default=.001, type=float) # try also .0003, .0001, go lower with lower batch size

    args, _ = _parse_args(config_parser, parser)

    if args.nhid is None:
        args.nhid = 2*args.emsize

    prior = args.__dict__.pop('prior')

    if prior == 'gp':
        prior = priors.fast_gp.DataLoader
    elif prior == 'ridge':
        prior = priors.ridge.DataLoader
    elif prior == 'stroke':
        prior = priors.stroke.DataLoader
    elif prior == 'mix_gp':
        prior = priors.fast_gp_mix.DataLoader
    else:
        raise NotImplementedError(f'Prior == {prior}.')

    loss_function = args.__dict__.pop('loss_function')

    criterion = nn.GaussianNLLLoss(reduction='none', full=True)
    classificiation_criterion = nn.CrossEntropyLoss(reduction='none')
    num_buckets = args.__dict__.pop('num_buckets')
    max_y = args.__dict__.pop('max_y')
    min_y = args.__dict__.pop('min_y')
    # criterion = nn.MSELoss(reduction='none')

    def get_y_sample():
        dl = prior(num_steps=1, batch_size=args.batch_size * args.steps_per_epoch, seq_len=args.bptt, device=device,
                   **args.extra_prior_kwargs_dict)
        y_sample = next(iter(dl))[-1]
        print(f'Creating Bar distribution with borders from y sample of size {y_sample.numel()}')
        return y_sample

    if loss_function == 'ce':
        criterion = nn.CrossEntropyLoss(reduction='none')
    elif loss_function == 'gaussnll':
        criterion = nn.GaussianNLLLoss(reduction='none', full=True)
    elif loss_function == 'mse':
        criterion = nn.MSELoss(reduction='none')
    elif loss_function == 'barnll':
        criterion = BarDistribution(borders=get_bucket_limits(num_buckets, full_range=(min_y,max_y)))
    elif loss_function == 'adaptivebarnll':
        borders = get_bucket_limits(num_buckets, ys=get_y_sample(), full_range=(min_y,max_y))
        criterion = BarDistribution(borders=borders)
    elif loss_function == 'adaptivefullsupportbarnll':
        assert min_y is None and max_y is None, "Please do not specify `min_y` and `max_y` with `unboundedadaptivebarnll`."
        borders = get_bucket_limits(num_buckets, ys=get_y_sample())
        criterion = FullSupportBarDistribution(borders=borders)
    else:
        raise NotImplementedError(f'loss_function == {loss_function}.')



    encoder = args.__dict__.pop('encoder')
    y_encoder = args.__dict__.pop('y_encoder')

    def get_encoder_generator(encoder):
        if encoder == 'linear':
            encoder_generator = encoders.Linear
        elif encoder == 'mlp':
            encoder_generator = encoders.MLP
        elif encoder == 'positional':
            encoder_generator = encoders.Positional
        else:
            raise NotImplementedError(f'A {encoder} encoder is not valid.')
        return encoder_generator

    encoder_generator = get_encoder_generator(encoder)
    y_encoder_generator = get_encoder_generator(y_encoder)

    pos_encoder = args.__dict__.pop('pos_encoder')

    if pos_encoder == 'none':
        pos_encoder_generator = None
    elif pos_encoder == 'sinus':
        pos_encoder_generator = positional_encodings.PositionalEncoding
    elif pos_encoder == 'learned':
        pos_encoder_generator = positional_encodings.LearnedPositionalEncoding
    elif pos_encoder == 'paired_scrambled_learned':
        pos_encoder_generator = positional_encodings.PairedScrambledPositionalEncodings
    else:
        raise NotImplementedError(f'pos_encoer == {pos_encoder} is not valid.')

    permutation_invariant_max_eval_pos = args.__dict__.pop('permutation_invariant_max_eval_pos')
    permutation_invariant_sampling = args.__dict__.pop('permutation_invariant_sampling')
    if permutation_invariant_max_eval_pos is not None:
        if permutation_invariant_sampling == 'weighted':
            get_sampler = get_weighted_single_eval_pos_sampler
        elif permutation_invariant_sampling == 'uniform':
            get_sampler = get_uniform_single_eval_pos_sampler
        else:
            raise ValueError()
        args.__dict__['single_eval_pos_gen'] = get_sampler(permutation_invariant_max_eval_pos)


    print("ARGS for `train`:", args.__dict__)

    train(prior, criterion, encoder_generator,
          y_encoder_generator=y_encoder_generator, pos_encoder_generator=pos_encoder_generator,
          **args.__dict__)