Flexpert-Design/train.py

import os, sys, warnings, argparse, math, tqdm, datetime
import pytorch_lightning as pl
import torch
from pytorch_lightning.trainer import Trainer
import pytorch_lightning.callbacks as plc
import pytorch_lightning.loggers as plog
from model_interface import MInterface
from data_interface import DInterface
from src.tools.logger import SetupCallback, BackupCodeCallback
from shutil import ignore_patterns

import wandb
warnings.filterwarnings("ignore")

def create_parser():
    parser = argparse.ArgumentParser()
    # Set-up parameters
    parser.add_argument('--res_dir', default='./train/results', type=str)
    parser.add_argument('--ex_name', default='debug', type=str)
    parser.add_argument('--check_val_every_n_epoch', default=1, type=int)
    parser.add_argument('--stage', default='fit', type=str) #'fit', 'test' or 'predict'
    parser.add_argument('--val_check_interval', default=0.5, type=float, help='Validation check interval')
    
    parser.add_argument('--dataset', default='FLEX_CATH4.3') # AF2DB_dataset, CATH_dataset
    parser.add_argument('--model_name', default='ProteinMPNN', choices=['StructGNN', 'GraphTrans', 'GVP', 'GCA', 'AlphaDesign', 'ESMIF', 'PiFold', 'ProteinMPNN', 'KWDesign', 'E3PiFold'])
    parser.add_argument('--lr', default=4e-4, type=float, help='Learning rate')
    parser.add_argument('--lr_scheduler', default='onecycle')
    parser.add_argument('--offline', default=1, type=int)
    parser.add_argument('--seed', default=111, type=int)
    
    # dataset parameters
    parser.add_argument('--batch_size', default=32, type=int)
    parser.add_argument('--num_workers', default=12, type=int)
    parser.add_argument('--pad', default=1024, type=int)
    parser.add_argument('--min_length', default=40, type=int)
    parser.add_argument('--data_root', default='../data/')
    parser.add_argument('--infer_path', default='', type=str)
    
    # Training parameters
    parser.add_argument('--epoch', default=10, type=int, help='end epoch')
    parser.add_argument('--augment_eps', default=0.0, type=float, help='noise level')
    parser.add_argument('--gpus', default=1, type=int, help='how many GPUs to train on')
    parser.add_argument('--weight_decay', default=0.0, type=float, help='Weight decay for optimizer')

    # Eval parameters
    parser.add_argument('--eval_sequences_sampled', default=1, type=int, help='How many sequences to sample in evaluation.')
    parser.add_argument('--eval_sequences_temperature', default=0, type=float, help='What temperature to use for the sampling in evaluation.')
    parser.add_argument('--eval_output_dir', default=None, type=str, help='Where to save the evaluation output.')

    # Model parameters
    parser.add_argument('--use_dist', default=1, type=int)
    parser.add_argument('--use_product', default=0, type=int)
    parser.add_argument('--use_pmpnn_checkpoint', type=int, help='By 1 or 0 decide whether to start with pretrained ProteinMPNN.')

    # Dynamics aware parameters
    parser.add_argument('--use_dynamics', default=0, type=int)
    parser.add_argument('--flex_loss_coeff', default=0.5, type=float)
    parser.add_argument('--get_gt_flex_onthefly', default=0, type=int, help='Flag to get ground truth flexibility on-the-fly (with subsequent caching)')
    parser.add_argument('--init_flex_features', default=1, type=int, help="Set to 0 if no flexibility information should be passed on input to the node features h_V")
    parser.add_argument('--loss_fn', default='MSE', type=str, help= 'Define what loss to use. Choose MSE, L1 or DPO.')
    parser.add_argument('--grad_normalization', default=0, type=int, help="Set to 0 if the gradients of the seq and flex losses should not be normalized.")
    parser.add_argument('--test_engineering', default=0, type=int, help="In this main.py should be set to 0 to not overwrite the training dataset.")
    
    args = parser.parse_args()
    return args




def load_callbacks(args):
    callbacks = []
    
    logdir = str(os.path.join(args.res_dir, args.ex_name))
    
    ckptdir = os.path.join(logdir, "checkpoints")
    
    callbacks.append(BackupCodeCallback(os.path.dirname(args.res_dir),logdir, ignore_patterns=ignore_patterns('results*', 'pdb*', 'metadata*', 'vq_dataset*')))
    

    metric = "recovery"
    sv_filename = 'best-{epoch:02d}-{recovery:.3f}'
    callbacks.append(plc.ModelCheckpoint(
        monitor=metric,
        filename=sv_filename,
        save_top_k=15,
        mode='max',
        save_last=True,
        dirpath = ckptdir,
        verbose = True,
        every_n_epochs = args.check_val_every_n_epoch,
    ))

    
    now = datetime.datetime.now().strftime("%m-%dT%H-%M-%S")
    cfgdir = os.path.join(logdir, "configs")
    callbacks.append(
        SetupCallback(
                now = now,
                logdir = logdir,
                ckptdir = ckptdir,
                cfgdir = cfgdir,
                config = args.__dict__,
                argv_content = sys.argv + ["gpus: {}".format(args.gpus)],)
    )
    
    
    if args.lr_scheduler:
        callbacks.append(plc.LearningRateMonitor(
            logging_interval=None))
    return callbacks



if __name__ == "__main__":
    
    args = create_parser()
    if args.stage == 'predict':
        args.batch_size = 1
        print('In the predict stage, defaulting batch size to 1.')


    if (len(args.infer_path) > 0 or args.dataset=='PDBInference') and (len(args.infer_path) == 0 or args.dataset!='PDBInference'):
        raise ValueError("You should only use --infer_path with --dataset 'PDBInference' and vice versa.")

    pl.seed_everything(args.seed)

    data_module = DInterface(**vars(args))

    data_module.setup(stage=args.stage) #here is the cache_data called

    gpu_count = args.gpus #torch.cuda.device_count()
    if args.stage == 'fit':
        args.steps_per_epoch = math.ceil(len(data_module.trainset)/args.batch_size/gpu_count)
        print(f"steps_per_epoch {args.steps_per_epoch},  gpu_count {gpu_count}, batch_size{args.batch_size}")
        

    model = MInterface(**vars(args))

    trainer_config = {
        'devices': args.gpus,
        'max_epochs': args.epoch,
        'num_nodes': 1,
        "strategy": 'ddp',
        "precision": '32',
        'accelerator': 'gpu',
        'callbacks': load_callbacks(args),
        'logger': plog.WandbLogger(
                    project = 'ICLR2025',
                    name=args.ex_name,
                    save_dir=str(os.path.join(args.res_dir, args.ex_name)),
                    offline = args.offline,
                    id = "_".join(args.ex_name.split("/")),
                    entity = "koubic"),
        'val_check_interval': args.val_check_interval,
        'check_val_every_n_epoch': args.check_val_every_n_epoch
    }

    trainer = Trainer(**trainer_config)

    if args.stage =='fit':
        trainer.fit(model, data_module)
    elif args.stage == 'test':
        test_out = trainer.test(model,data_module)
    elif args.stage == 'eval':
        from transformers import AutoTokenizer
        tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t33_650M_UR50D", cache_dir='./cache_dir/')
        
        predictions = trainer.predict(model, data_module) #Just 1015 proteins loaded since they are required to be shorther than 500 residues
        out_dict = {}
        for pred in tqdm.tqdm(predictions):
            logprobs = pred['log_probs']
            for i in range(args.eval_sequences_sampled):

                if args.eval_sequences_temperature > 0 or args.eval_sequences_sampled > 1:
                    raise NotImplementedError('Sampling with temperature is not implemented yet.') #TODO!!!
                else:
                    AA_indices = logprobs.argmax(dim=-1, keepdim=False)
                    
                    decoded_seqs = tokenizer.batch_decode(AA_indices,skip_special_tokens=True)

                    for title, seq, mask in zip(pred['title'],decoded_seqs, pred['mask']):
                        _seq = [letter for letter,cond in zip(seq.split(' '),mask) if cond]
                        seq_cat = ''.join(_seq)
                        out_dict[title] = seq_cat
        with open(os.path.join(args.eval_output_dir,'inverse_folded_ATLAS.fasta'), 'w') as f:
            for pdb_name, seq in out_dict.items(): 
                f.write(f">{pdb_name}\n{seq}\n")

    elif args.stage == 'predict':
        predictions = trainer.predict(model, data_module)
        predictions_cuda = []
        for pred in predictions:
            _pred_cuda = {}
            for k, v in pred.items():
                if isinstance(v, torch.Tensor):
                    _pred_cuda[k] = v.to(torch.device('cuda'))
            _pred_cuda['batch'] = {k: pred['batch'][k].to(torch.device('cuda')) for k in ('X', 'S', 'mask', 'chain_M', 'chain_M_pos', 'residue_idx', 'chain_encoding_all')}
            predictions_cuda.append(_pred_cuda)

        seriazable_predictions = []
        for pred in predictions:
            ser_pred = {}
            for k, v in pred.items():
                if isinstance(v, torch.Tensor):
                    ser_pred[k] = v.cpu().numpy().tolist()
                else:
                    ser_pred[k] = v
            seriazable_predictions.append(ser_pred)

        import json
        with open(f'{args.infer_path}/predictions.json', 'w') as f:
            json.dump(seriazable_predictions, f)