| """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) |
| |
| 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') |
| |
| 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() |
|
|
