train_mono.py

import argparse
import os
import random
import time
import warnings

import esm
import numpy as np
import torch
import torch.distributed as dist
from torch.utils.data import DataLoader, DistributedSampler

from datasets.protdataset import ProtSeqDETRDataset
from engine import evaluate_mono_zero_level, evaluate_mono_enzyme_level, train_one_epoch
from models import build_model
import util.misc as utils
from util.clean import get_ec_id_dict, get_id_seq_dict

warnings.filterwarnings("ignore")

def get_dist_args():
  envvars = [
    "WORLD_SIZE",
    "RANK",
    "LOCAL_RANK",
    "NODE_RANK",
    "NODE_COUNT",
    "HOSTNAME",
    "MASTER_ADDR",
    "MASTER_PORT",
    "NCCL_SOCKET_IFNAME",
    "OMPI_COMM_WORLD_RANK",
    "OMPI_COMM_WORLD_SIZE",
    "OMPI_COMM_WORLD_LOCAL_RANK",
    "AZ_BATCHAI_MPI_MASTER_NODE",
  ]
  args = dict(gpus_per_node=torch.cuda.device_count())
  missing = []
  for var in envvars:
    if var in os.environ:
      args[var] = os.environ.get(var)
      try:
        args[var] = int(args[var])
      except ValueError:
        pass
    else:
      missing.append(var)

  print(f"II Args: {args}")
  if missing:
    print(f"II Environment variables not set: {', '.join(missing)}.")
    
  return args

def get_args_parser():
    parser = argparse.ArgumentParser('Set transformer detector', add_help=False)
    # training args
    parser.add_argument('--lr', default=1e-4, type=float)
    parser.add_argument('--lr_backbone', default=1e-5, type=float)
    parser.add_argument('--batch_size', default=8, type=int)
    parser.add_argument('--weight_decay', default=1e-4, type=float)
    parser.add_argument('--epochs', default=30, type=int)
    parser.add_argument('--lr_drop', default=200, type=int)
    parser.add_argument('--clip_max_norm', default=0.1, type=float,
                        help='gradient clipping max norm')
    parser.add_argument('--model_name', type=str, default='ProtDETR_ECPred40')
    # Backbone
    parser.add_argument('--position_embedding', default='sine', type=str, choices=('sine', 'learned'),
                        help="Type of positional embedding to use on top of the image features")
    # Transformer
    parser.add_argument('--enc_layers', default=3, type=int,
                        help="Number of encoding layers in the transformer")
    parser.add_argument('--dec_layers', default=3, type=int,
                        help="Number of decoding layers in the transformer")
    parser.add_argument('--dim_feedforward', default=2048, type=int,
                        help="Intermediate size of the feedforward layers in the transformer blocks")
    parser.add_argument('--hidden_dim', default=256, type=int,
                        help="Size of the embeddings (dimension of the transformer)")
    parser.add_argument('--dropout', default=0.1, type=float,
                        help="Dropout applied in the transformer")
    parser.add_argument('--nheads', default=4, type=int,
                        help="Number of attention heads inside the transformer's attentions")
    parser.add_argument('--num_queries', default=1, type=int,
                        help="Number of query slots")
    parser.add_argument('--pre_norm', action='store_true')

    # Loss
    parser.add_argument('--no_aux_loss', dest='aux_loss', action='store_false',
                        help="Disables auxiliary decoding losses (loss at each layer)")
    # Matcher
    parser.add_argument('--set_cost_class', default=1, type=float,
                        help="Class coefficient in the matching cost")
    # Loss coefficients
    parser.add_argument('--eos_coef', default=0.0, type=float,
                        help="Relative classification weight of the no-enzyme class")
    # dataset parameters
    parser.add_argument('--start_epoch', default=0, type=int, metavar='N',
                        help='start epoch')
    parser.add_argument('--device', default='cuda',
                        help='device to use for training / testing')
    parser.add_argument('--seed', default=42, type=int)
    parser.add_argument('--num_workers', default=2, type=int)
    parser.add_argument('--esm_layer', default=32, type=int)
    return parser


def main(args):
    # training args
    args = get_args_parser().parse_args()
    print(args)
    # ddp args
    dist_args = get_dist_args()
    if 'RANK' in os.environ and 'WORLD_SIZE' in os.environ:
        args.distributed = True
    else:
        args.distributed = False
    if args.distributed:
        master_uri = "tcp://%s:%s" % (dist_args.get("MASTER_ADDR"), dist_args.get("MASTER_PORT"))
        os.environ["NCCL_DEBUG"] = "WARN"
        node_rank = dist_args.get("NODE_RANK")
        gpus_per_node = torch.cuda.device_count()
        world_size = dist_args.get("WORLD_SIZE")
        gpu_rank = dist_args.get("LOCAL_RANK")
        node_rank = 0 if node_rank is None else node_rank
        global_rank = node_rank * gpus_per_node + gpu_rank
        dist.init_process_group(
            backend="nccl", init_method=master_uri, world_size=world_size, rank=global_rank
        )
        torch.cuda.set_device(gpu_rank)
        device = torch.device("cuda", gpu_rank)
    else:
        device = torch.device(args.device)
    # fix the seed for reproducibility
    seed = args.seed + utils.get_rank()
    torch.manual_seed(seed)
    np.random.seed(seed)
    random.seed(seed)

    train_data_pth = f"./data/mono_func/ECPred40_train.csv"
    id_ec_train, ec_id_train = get_ec_id_dict(train_data_pth)
    id_seq_train = get_id_seq_dict(train_data_pth)
    
    train_dataset = ProtSeqDETRDataset(id_ec_train, ec_id_train, id_seq_train, max_labes=args.num_queries, esm_layer=args.esm_layer)
    ec_to_label = train_dataset.ec_to_label
    label_to_ec = train_dataset.label_to_ec
    num_labels = len(ec_to_label)
    args.num_classes = num_labels

    # this valid datset removes those non-enzyme instances to evaluate the enzyme level(1, 2, 3, 4) following EnzBert
    # we select the best model based on the level 4 F1 score
    valid_data_path_remove_zero = f'./data/mono_func/ECPred40RemoveZero_valid.csv'
    id_ec_valid_remove_zero, ec_id_valid_remove_zero = get_ec_id_dict(valid_data_path_remove_zero)
    id_seq_valid_remove_zero = get_id_seq_dict(valid_data_path_remove_zero)
    valid_dataset_remove_zero = ProtSeqDETRDataset(id_ec_valid_remove_zero, ec_id_valid_remove_zero, id_seq_valid_remove_zero, max_labes=args.num_queries, esm_layer=args.esm_layer, ec_to_label=ec_to_label, label_to_ec=label_to_ec)

    # download esm model only once  
    if args.distributed:
        if utils.get_rank() == 0:
            esm_model, alphabet = esm.pretrained.esm1b_t33_650M_UR50S()
            torch.distributed.barrier()
        else:
            torch.distributed.barrier()
            esm_model, alphabet = esm.pretrained.esm1b_t33_650M_UR50S()
    else:
        esm_model, alphabet = esm.pretrained.esm1b_t33_650M_UR50S()
        
    esm_model.eval()
    esm_model.to(device)

    model, criterion = build_model(args, train_dataset.ec_weight)
    model.to(device)
    criterion.to(device)
    model_without_ddp = model
    if args.distributed:
        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[gpu_rank], output_device=gpu_rank,)
        model_without_ddp = model.module
    n_parameters = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print('number of params:', n_parameters)

    param_dicts = [
        {"params": [p for n, p in model_without_ddp.named_parameters() if "backbone" not in n and p.requires_grad]},
        {
            "params": [p for n, p in model_without_ddp.named_parameters() if "backbone" in n and p.requires_grad],
            "lr": args.lr_backbone,
        },
    ]
    optimizer = torch.optim.AdamW(param_dicts, lr=args.lr,
                                  weight_decay=args.weight_decay)
    lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, args.lr_drop)

    if args.distributed:
        sampler_train = DistributedSampler(train_dataset)
        train_loader = DataLoader(train_dataset, batch_size=args.batch_size, sampler=sampler_train, collate_fn=train_dataset.collate_fn, num_workers=args.num_workers)
    else:
        train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, collate_fn=train_dataset.collate_fn, num_workers=args.num_workers)
    
    valid_loader_remove_zero = DataLoader(valid_dataset_remove_zero, batch_size=args.batch_size, shuffle=False, collate_fn=valid_dataset_remove_zero.collate_fn, num_workers=args.num_workers)

    print("Start training")
    best_f1 = 0.0
    for epoch in range(args.start_epoch, args.epochs):
        epoch_start_time = time.time()
        if args.distributed:
            sampler_train.set_epoch(epoch)
        train_loss, esm_time = train_one_epoch(esm_model, alphabet, args.esm_layer,
            model, criterion, train_loader, optimizer, device, epoch, args.clip_max_norm)
        lr_scheduler.step()
        train_end_time = time.time()
        if utils.get_rank() == 0:
            print(f'Epoch {epoch}/{args.epochs} | Train Loss: {train_loss:.4f} | Time: {(train_end_time - epoch_start_time):.2f}s | ESM Time: {esm_time:.2f}s')

        if utils.get_rank() == 0:
            if (epoch + 1) % 1 == 0: # evaluate every epoch
                eval_start_time = time.time()
                f1 = evaluate_mono_enzyme_level(args.model_name, esm_model, alphabet, args.esm_layer, model, "ECPred40_valid_remove_zero", valid_loader_remove_zero, ec_to_label, label_to_ec, args.num_queries, device)
                eval_end_time = time.time()
                print(f"Evaluate time: {(eval_end_time - eval_start_time):.2f}s")
                if f1 > best_f1:
                    best_f1 = f1
                    print(f"Best F1: {best_f1:.4f} at epoch {epoch}")
                    if dist.is_initialized():
                        torch.save(model.module.state_dict(), f"./saved_models/{args.model_name}.pt")
                    else:
                        torch.save(model.state_dict(), f"./saved_models/{args.model_name}.pt")

if __name__ == '__main__':
    parser = argparse.ArgumentParser('ProtDETR mono-func training script', parents=[get_args_parser()])
    args = parser.parse_args()
    main(args)