Blame for Computational/package/train.py

"""Train MetaLATTE or Metalorian."""
import argparse
import csv
import json
from pathlib import Path


def arguments():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument('--stage', choices=['metalatte', 'metalorian'], required=True)
    p.add_argument('--output', type=Path, required=True)
    p.add_argument('--device', default='cuda:0')
    p.add_argument('--epochs', type=int, default=30)
    p.add_argument('--batch-size', type=int, default=2,
                   help='Sequence batch size for Metalorian; MetaLATTE uses its saved pre-batches')
    p.add_argument('--workers', type=int, default=0)
    p.add_argument('--threads', type=int, default=4)
    p.add_argument('--seed', type=int, default=42)
    p.add_argument('--lr', type=float, default=2e-4)
    p.add_argument('--init-checkpoint', type=Path,
                   help='MetaLATTE weights to initialize from; defaults to the supplied predictor for Metalorian')
    p.add_argument('--resume', type=Path, help='Resume a checkpoint for the selected stage')
    p.add_argument('--smoke-test', action='store_true', help='One update using at most two records and 64 tokens; saves no model')
    a = p.parse_args()
    if min(a.epochs, a.batch_size, a.threads) < 1 or a.workers < 0:
        p.error('epochs, batch-size and threads must be positive; workers cannot be negative')
    if a.smoke_test and a.resume:
        p.error('smoke-test starts a short test run; omit resume')
    return a


def class_weights():
    import numpy as np
    import torch
    from datasets import load_from_disk
    from metalorian.paths import TRAIN
    ds = load_from_disk(str(TRAIN))
    counts = np.zeros(15, dtype=np.float64)
    for row in ds.select_columns(['anchor_labels']):
        counts += np.asarray(row['anchor_labels']).sum(axis=0)
    weights = 1.0/(counts+1e-6)
    return torch.tensor(weights/weights.sum()*15, dtype=torch.float32)


def predictor(args, checkpoint=None, residue_encoder=False):
    import torch
    from transformers import EsmModel, EsmConfig
    from metalorian.paths import BACKBONE, PREDICTOR
    from metalorian.predictor import MultitaskProteinModel, ResidueEncoder
    metadata = torch.load(PREDICTOR, map_location='cpu', weights_only=False, mmap=True)
    hp = argparse.Namespace(**metadata['hyper_parameters'])
    hp.lr, hp.epochs = args.lr, args.epochs
    hp.output_model_path = str(args.output.resolve())
    del metadata
    esm = (EsmModel(EsmConfig.from_pretrained(str(BACKBONE), local_files_only=True))
           if checkpoint else EsmModel.from_pretrained(str(BACKBONE), local_files_only=True))
    cls = ResidueEncoder if residue_encoder else MultitaskProteinModel
    model = cls(hparams=hp, esm_model=esm, num_layers_to_finetune=2,
                input_dim=1280, output_dim=15, class_weights=class_weights())
    if checkpoint:
        state = torch.load(checkpoint, map_location='cpu', weights_only=False, mmap=True)
        model.load_state_dict(state['state_dict'], strict=True)
    return model


def smoke_group(row):
    result = {}
    for key, values in row.items():
        result[key] = values[:2] if key.endswith('labels') else [x[:64] for x in values[:2]]
    return result


def train_metalatte(args):
    import torch
    from datasets import load_from_disk
    from torch.utils.data import DataLoader
    from lightning.pytorch import Trainer
    from lightning.pytorch.callbacks import ModelCheckpoint, StochasticWeightAveraging
    from lightning.pytorch.loggers import CSVLogger
    from metalorian.paths import TRAIN, VALIDATION
    from metalorian.data import collate_fn_hard, collate_fn
    train, valid = load_from_disk(str(TRAIN)), load_from_disk(str(VALIDATION))
    if args.smoke_test:
        train, valid = [smoke_group(train[0])], [smoke_group(valid[0])]
    model = predictor(args, args.init_checkpoint or args.resume)
    train_loader = DataLoader(train, batch_size=1, shuffle=False, collate_fn=collate_fn_hard, num_workers=args.workers)
    valid_loader = DataLoader(valid, batch_size=1, shuffle=False, collate_fn=collate_fn, num_workers=args.workers)
    gpu = args.device.startswith('cuda')
    device_index = int(args.device.split(':')[1]) if ':' in args.device else 0
    callbacks = [] if args.smoke_test else [
        ModelCheckpoint(dirpath=args.output/'checkpoints', save_last=True, save_top_k=1,
                        monitor='val_macro_f1', mode='max', filename='metalatte-{epoch:02d}'),
        StochasticWeightAveraging(swa_lrs=1e-2)]
    trainer = Trainer(max_epochs=args.epochs, accelerator='gpu' if gpu else 'cpu',
                      devices=[device_index] if gpu else 1, precision='bf16-mixed' if gpu else '32-true',
                      logger=CSVLogger(str(args.output), name='metrics'), callbacks=callbacks,
                      gradient_clip_val=0.5, enable_checkpointing=not args.smoke_test,
                      enable_model_summary=False, fast_dev_run=args.smoke_test,
                      default_root_dir=str(args.output))
    trainer.fit(model, train_dataloaders=train_loader, val_dataloaders=valid_loader,
                ckpt_path=str(args.resume) if args.resume else None)
    return {'stage': 'metalatte', 'status': 'completed', 'steps': int(trainer.global_step),
            'smoke_test': args.smoke_test}


def train_metalorian(args):
    import torch
    import numpy as np
    from datasets import load_from_disk
    from torch.utils.data import DataLoader, DistributedSampler
    from types import SimpleNamespace
    from metalorian.paths import ROOT, PREDICTOR
    from metalorian.continuous import DiffusionProteinModel, GaussianDiffusionTrainer
    from metalorian.discrete import MultinomialDiffusion
    from metalorian.unet import tabularUnet
    from metalorian.objectives import training_with_esm, make_negative_condition_esm, warmup_lr

    flags = SimpleNamespace(T=50, beta_1=1e-5, beta_T=0.02, input_size=1280,
        cond_size=15, output_size=1280, encoder_dim=[64,128,256], nf=16, activation='relu',
        hidden_size=768, num_layers=12, num_heads=12, intermediate_size=3072,
        max_seq_length=1280, lambda_con=0.2, lambda_dis=0.2)
    base = predictor(args, args.init_checkpoint or PREDICTOR, residue_encoder=True)
    con = DiffusionProteinModel(base, flags).to(args.device)
    flags.input_size, flags.cond_size, flags.output_size, flags.nf = 15, 1280, 15, 64
    dis = tabularUnet(flags).to(args.device)
    con_trainer = GaussianDiffusionTrainer(con, flags.beta_1, flags.beta_T, flags.T).to(args.device)
    data = load_from_disk(str(ROOT/'data/metalorian/train'))
    if args.smoke_test:
        data = [{k:(v[:64] if k in {'input_ids','attention_mask'} else v) for k,v in data[i].items()}
                for i in range(min(2, len(data)))]
    data = TokenRows(data)
    # Preserve the training sampler order: seed 0, epoch 0.
    order = DistributedSampler(data, num_replicas=1, rank=0, shuffle=True, seed=0)
    loader = DataLoader(data, batch_size=args.batch_size, sampler=order, num_workers=args.workers)
    dis_trainer = MultinomialDiffusion(np.array([15]), (len(data),15), dis, flags,
        timesteps=flags.T, loss_type='vb_stochastic').to(args.device)
    optim_con = torch.optim.Adam(con.parameters(), lr=args.lr)
    optim_dis = torch.optim.Adam(dis.parameters(), lr=args.lr)
    sched_con = torch.optim.lr_scheduler.LambdaLR(optim_con, warmup_lr)
    sched_dis = torch.optim.lr_scheduler.LambdaLR(optim_dis, warmup_lr)
    start_epoch, step = 0, 0
    if args.resume:
        ckpt = torch.load(args.resume, map_location='cpu', weights_only=False, mmap=True)
        strip = lambda state: {k.removeprefix('module.'):v for k,v in state.items()}
        con.load_state_dict(strip(ckpt['model_con']), strict=True)
        dis.load_state_dict(strip(ckpt['model_dis']), strict=True)
        optim_con.load_state_dict(ckpt['optim_con']); optim_dis.load_state_dict(ckpt['optim_dis'])
        sched_con.load_state_dict(ckpt['sched_con']); sched_dis.load_state_dict(ckpt['sched_dis'])
        start_epoch, step = int(ckpt.get('epoch',0)), int(ckpt.get('step',0))+1
        del ckpt
    con_params = [p for p in con.parameters() if p.requires_grad]
    dis_params = [p for p in dis.parameters() if p.requires_grad]
    transform = SimpleNamespace(output_info=[(15,'softmax')])
    weights = class_weights().to(args.device)
    with (args.output/'losses.csv').open('w', newline='') as handle:
        log = csv.writer(handle)
        log.writerow(['epoch','step','continuous_loss','continuous_contrastive','discrete_loss','discrete_contrastive'])
        for epoch in range(start_epoch, start_epoch+args.epochs):
            con.train(); dis.train()
            for batch in loader:
                x = {k:batch[k].to(args.device) for k in ['input_ids','attention_mask']}
                labels = batch['labels'].float().to(args.device)
                ns_con, ns_dis = make_negative_condition_esm(x, labels, base)
                values = training_with_esm(x, labels, con_trainer, dis_trainer, ns_con, ns_dis,
                                          transform, flags, None, weights)
                loss_con = values[0]+flags.lambda_con*values[1]
                loss_dis = values[2]+flags.lambda_dis*values[3]
                if not torch.isfinite(loss_con+loss_dis):
                    raise RuntimeError('Nonfinite training loss')
                # Compute both gradients before updating the shared forward graph.
                optim_con.zero_grad(set_to_none=True); optim_dis.zero_grad(set_to_none=True)
                gc = torch.autograd.grad(loss_con, con_params, retain_graph=True, allow_unused=True)
                gd = torch.autograd.grad(loss_dis, dis_params, allow_unused=True)
                for p,g in zip(con_params,gc): p.grad=g
                for p,g in zip(dis_params,gd): p.grad=g
                torch.nn.utils.clip_grad_norm_(con_params, 1.0)
                torch.nn.utils.clip_grad_norm_(dis_params, 1.0)
                optim_con.step(); sched_con.step()
                optim_dis.step(); sched_dis.step()
                numbers = [float(v.detach()) if torch.is_tensor(v) else float(v) for v in values]
                log.writerow([epoch,step,*numbers]); handle.flush()
                print(f'epoch={epoch} step={step} continuous={float(loss_con.detach()):.6f} discrete={float(loss_dis.detach()):.6f}', flush=True)
                step += 1
                if args.smoke_test:
                    return {'stage':'metalorian','status':'completed','steps':1,'smoke_test':True,'losses':numbers}
            torch.save(dict(model_con=con.state_dict(), model_dis=dis.state_dict(),
                optim_con=optim_con.state_dict(), optim_dis=optim_dis.state_dict(),
                sched_con=sched_con.state_dict(), sched_dis=sched_dis.state_dict(),
                step=step-1, epoch=epoch+1), args.output/'metalorian_latest.pt')
    return {'stage':'metalorian','status':'completed','steps':step,'smoke_test':False}


class TokenRows:
    def __init__(self, data): self.data=data
    def __len__(self): return len(self.data)
    def __getitem__(self, index):
        import torch
        row=self.data[index]
        return {k:torch.tensor(row[k]) for k in ['input_ids','attention_mask','labels']}


def main():
    args = arguments()
    from metalorian.paths import configure
    configure()
    import torch
    from lightning.pytorch import seed_everything
    seed_everything(args.seed, workers=True)
    torch.set_num_threads(args.threads)
    args.output.mkdir(parents=True, exist_ok=True)
    (args.output/'settings.json').write_text(json.dumps({k:str(v) if isinstance(v,Path) else v
                                                       for k,v in vars(args).items()}, indent=2)+'\n')
    result = train_metalatte(args) if args.stage=='metalatte' else train_metalorian(args)
    (args.output/'result.json').write_text(json.dumps(result, indent=2)+'\n')
    print(json.dumps(result, indent=2))


if __name__ == '__main__':
    main()