train.py

import argparse
import os
import shutil

import numpy as np
import torch
# import torch.utils.tensorboard
from sklearn.metrics import roc_auc_score
from torch.nn.utils import clip_grad_norm_
from torch_geometric.loader import DataLoader
from torch_geometric.transforms import Compose
from tqdm.auto import tqdm

import utils.misc as misc
import utils.train as utils_train
import utils.transforms as trans
from datasets import get_dataset
from datasets.pl_data import FOLLOW_BATCH
from models.molopt_score_model import ScorePosNet3D
from models.geodesic_phi import GeodesicPhiNet

from graphbap.bapnet import BAPNet

def get_auroc(y_true, y_pred, feat_mode):
    y_true = np.array(y_true)
    y_pred = np.array(y_pred)
    avg_auroc = 0.
    possible_classes = set(y_true)
    for c in possible_classes:
        auroc = roc_auc_score(y_true == c, y_pred[:, c])
        avg_auroc += auroc * np.sum(y_true == c)
        mapping = {
            'basic': trans.MAP_INDEX_TO_ATOM_TYPE_ONLY,
            'add_aromatic': trans.MAP_INDEX_TO_ATOM_TYPE_AROMATIC,
            'full': trans.MAP_INDEX_TO_ATOM_TYPE_FULL
        }
        print(f'atom: {mapping[feat_mode][c]} \t auc roc: {auroc:.4f}')
    return avg_auroc / len(y_true)


if __name__ == '__main__':
    root_dir = 'C:/Users/MinhTQ/FlowMatching-SBDD'

    parser = argparse.ArgumentParser()
    parser.add_argument('--config', type=str, default=root_dir + '/configs/training.yml')
    parser.add_argument('--device', type=str, default='cuda')
    parser.add_argument('--logdir', type=str, default=root_dir + '/logs')
    parser.add_argument('--tag', type=str, default='')
    parser.add_argument('--train_report_iter', type=int, default=200)
    # Resume để train tiếp khi Kaggle ngắt phiên (giới hạn 12h)
    parser.add_argument('--resume', type=str, default=None,
                        help='Đường dẫn checkpoint (.pt, vd last.pt) để train tiếp từ iteration đã lưu')
    parser.add_argument('--ckpt_freq', type=int, default=2000,
                        help='Tần suất (số iter) lưu last.pt để luôn resume được dù bị ngắt giữa chừng')
    args = parser.parse_args()

    # Load configs
    config = misc.load_config(args.config)
    config_name = os.path.basename(args.config)[:os.path.basename(args.config).rfind('.')]
    misc.seed_all(config.train.seed)

    # Logging
    log_dir = misc.get_new_log_dir(args.logdir, prefix=config_name, tag=args.tag)
    ckpt_dir = os.path.join(log_dir, 'checkpoints')
    os.makedirs(ckpt_dir, exist_ok=True)
    vis_dir = os.path.join(log_dir, 'vis')
    os.makedirs(vis_dir, exist_ok=True)
    logger = misc.get_logger('train', log_dir)

    logger.info(args)
    logger.info(config)
    shutil.copyfile(args.config, os.path.join(log_dir, os.path.basename(args.config)))
    shutil.copytree(root_dir + '/models', os.path.join(log_dir, 'models'))

    # Transforms & vetorizer
    protein_featurizer = trans.FeaturizeProteinAtom()
    ligand_featurizer = trans.FeaturizeLigandAtom(config.data.transform.ligand_atom_mode)
    transform_list = [
        protein_featurizer,
        ligand_featurizer,
        trans.FeaturizeLigandBond(),
    ]
    if config.data.transform.random_rot:
        transform_list.append(trans.RandomRotation())
    transform = Compose(transform_list)

    # Datasets and loaders
    logger.info('Loading dataset...')
    dataset, subsets = get_dataset(
        config=config.data,
        transform=transform,
    )
    train_set, val_set = subsets['train'], subsets['test']
    logger.info(f'Training: {len(train_set)} Validation: {len(val_set)}')

    collate_exclude_keys = ['ligand_nbh_list']
    train_iterator = utils_train.inf_iterator(DataLoader(
        train_set,
        batch_size=config.train.batch_size,
        shuffle=True,
        num_workers=config.train.num_workers,
        follow_batch=FOLLOW_BATCH,
        exclude_keys=collate_exclude_keys
    ))
    val_loader = DataLoader(val_set, config.train.batch_size, shuffle=False,
                            follow_batch=FOLLOW_BATCH, exclude_keys=collate_exclude_keys)
    # Model
    logger.info('Building model...')

    # IPNet model
    net_cond = BAPNet(ckpt_path=config.net_cond.ckpt_path, hidden_nf=config.net_cond.hidden_dim).to(args.device)
    for param in net_cond.parameters():
        param.requires_grad = False  # Freeze

    # Mô hình 
    model = ScorePosNet3D(
        # net_cond,
        config.model,
        protein_atom_feature_dim=protein_featurizer.feature_dim,
        ligand_atom_feature_dim=ligand_featurizer.feature_dim
    ).to(args.device)

    print(f'protein feature dim: {protein_featurizer.feature_dim} ligand feature dim: {ligand_featurizer.feature_dim}')
    logger.info(f'# trainable parameters: {misc.count_parameters(model) / 1e6:.4f} M')

    # Optimizer and scheduler
    optimizer = utils_train.get_optimizer(config.train.optimizer, model)
    scheduler = utils_train.get_scheduler(config.train.scheduler, optimizer)

    use_fm = getattr(config.model, 'use_flow_matching', False)
    fm_variant = getattr(config.model, 'fm_variant', 'pure')
    # AUX CE (chỉ FM): 0 = tắt, >0 = bật với trọng số đó. Đặt ở config.train.aux_ce_weight.
    aux_ce_weight = getattr(config.train, 'aux_ce_weight', 1.0)

    # MFM stage 2: nạp phi (đã train ở stage 1) và ĐÔNG BĂNG để dựng velocity_target.
    # Chỉ cần khi fm_variant=='geodesic'; các variant khác không đụng tới phi.
    phi_net = None
    if use_fm and fm_variant == 'geodesic':
        phi_ckpt_path = getattr(config.model, 'phi_ckpt', None)
        assert phi_ckpt_path is not None, \
            "fm_variant='geodesic' cần config.model.phi_ckpt trỏ tới phi_best.pt"
        phi_ckpt = torch.load(phi_ckpt_path, map_location=args.device, weights_only=False)
        phi_cfg = phi_ckpt['config'].phi   # tái dựng đúng siêu tham số đã train
        phi_net = GeodesicPhiNet(
            ligand_atom_feature_dim=ligand_featurizer.feature_dim,
            protein_atom_feature_dim=protein_featurizer.feature_dim,
            hidden_dim=phi_cfg.hidden_dim,
            num_layers=phi_cfg.num_layers,
            knn=phi_cfg.knn,
            num_r_gaussian=phi_cfg.num_r_gaussian,
            edge_feat_dim=phi_cfg.edge_feat_dim,
        ).to(args.device)
        phi_net.load_state_dict(phi_ckpt['model'])
        for p in phi_net.parameters():
            p.requires_grad = False   # đông băng tham số
        phi_net.eval()                # tắt dropout/cập-nhật buffer (norm=False nên không có BN)
        logger.info(f'Loaded frozen phi from {phi_ckpt_path} '
                    f'(val_loss={phi_ckpt.get("val_loss", "?")})')

    def train(it):
        model.train()
        optimizer.zero_grad()
        for _ in range(config.train.n_acc_batch):
            batch = next(train_iterator).to(args.device)

            protein_noise = torch.randn_like(batch.protein_pos) * config.train.pos_noise_std
            gt_protein_pos = batch.protein_pos + protein_noise

            if use_fm:
                results = model.get_flow_matching_loss(
                    net_cond=net_cond,
                    protein_pos=gt_protein_pos,
                    protein_v=batch.protein_atom_feature.float(),
                    batch_protein=batch.protein_element_batch,
                    ligand_pos=batch.ligand_pos,
                    ligand_v=batch.ligand_atom_feature_full,
                    batch_ligand=batch.ligand_element_batch,
                    fm_variant=fm_variant,
                    phi_net=phi_net,
                    aux_ce_weight=aux_ce_weight
                )
            else:
                results = model.get_diffusion_loss(
                    net_cond=net_cond,
                    protein_pos=gt_protein_pos,
                    protein_v=batch.protein_atom_feature.float(),
                    batch_protein=batch.protein_element_batch,
                    ligand_pos=batch.ligand_pos,
                    ligand_v=batch.ligand_atom_feature_full,
                    batch_ligand=batch.ligand_element_batch
                )
            loss, loss_pos, loss_v = results['loss'], results['loss_pos'], results['loss_v']
            loss = loss / config.train.n_acc_batch
            loss.backward()
        orig_grad_norm = clip_grad_norm_(model.parameters(), config.train.max_grad_norm)
        optimizer.step()

        if it % args.train_report_iter == 0:
            logger.info(
                '[Train] Iter %d | Loss %.6f (pos %.6f | v %.6f) | Lr: %.6f | Grad Norm: %.6f' % (
                    it, loss, loss_pos, loss_v, optimizer.param_groups[0]['lr'], orig_grad_norm
                )
            )


    def validate(it):
        # fix time steps
        sum_loss, sum_loss_pos, sum_loss_v, sum_n = 0, 0, 0, 0
        sum_loss_bond, sum_loss_non_bond = 0, 0
        all_pred_v, all_true_v = [], []
        pred_v_by_t, true_v_by_t = {}, {}   # AUROC theo TỪNG mức nhiễu t (chẩn đoán type head hỏng ở t nào)
        pos_sum_by_t = {}                    # Đo 2: loss_pos theo từng t (vị trí học thật hay chỉ copy?)
        all_pred_bond_type, all_gt_bond_type = [], []
        with torch.no_grad():
            model.eval()
            for batch in tqdm(val_loader, desc='Validate'):
                batch = batch.to(args.device)
                batch_size = batch.num_graphs
                t_loss, t_loss_pos, t_loss_v = [], [], []
                val_t_seq = (np.linspace(0, 1, 10) * (model.num_timesteps - 1)).astype(int) if use_fm \
                    else np.linspace(0, model.num_timesteps - 1, 10).astype(int)
                for t in val_t_seq:
                    time_step = torch.tensor([t] * batch_size).to(args.device)
                    if use_fm:
                        results = model.get_flow_matching_loss(
                            net_cond=net_cond,
                            protein_pos=batch.protein_pos,
                            protein_v=batch.protein_atom_feature.float(),
                            batch_protein=batch.protein_element_batch,
                            ligand_pos=batch.ligand_pos,
                            ligand_v=batch.ligand_atom_feature_full,
                            batch_ligand=batch.ligand_element_batch,
                            fm_variant=fm_variant,
                            time_step=time_step,
                            phi_net=phi_net,
                            aux_ce_weight=aux_ce_weight
                        )
                    else:
                        results = model.get_diffusion_loss(
                            net_cond=net_cond,
                            protein_pos=batch.protein_pos,
                            protein_v=batch.protein_atom_feature.float(),
                            batch_protein=batch.protein_element_batch,
                            ligand_pos=batch.ligand_pos,
                            ligand_v=batch.ligand_atom_feature_full,
                            batch_ligand=batch.ligand_element_batch,
                            time_step=time_step
                        )
                    loss, loss_pos, loss_v = results['loss'], results['loss_pos'], results['loss_v']

                    sum_loss += float(loss) * batch_size
                    sum_loss_pos += float(loss_pos) * batch_size
                    sum_loss_v += float(loss_v) * batch_size
                    sum_n += batch_size
                    all_pred_v.append(results['ligand_v_recon'].detach().cpu().numpy())
                    all_true_v.append(batch.ligand_atom_feature_full.detach().cpu().numpy())
                    ti = int(t)
                    pred_v_by_t.setdefault(ti, []).append(results['ligand_v_recon'].detach().cpu().numpy())
                    true_v_by_t.setdefault(ti, []).append(batch.ligand_atom_feature_full.detach().cpu().numpy())
                    # Đo 2: loss_pos theo t
                    if ti not in pos_sum_by_t:
                        pos_sum_by_t[ti] = [0.0, 0]
                    pos_sum_by_t[ti][0] += float(loss_pos) * batch_size
                    pos_sum_by_t[ti][1] += batch_size

        avg_loss = sum_loss / sum_n
        avg_loss_pos = sum_loss_pos / sum_n
        avg_loss_v = sum_loss_v / sum_n
        atom_auroc = get_auroc(np.concatenate(all_true_v), np.concatenate(all_pred_v, axis=0),
                               feat_mode=config.data.transform.ligand_atom_mode)

        # === CHẨN ĐOÁN: AUROC theo từng mức nhiễu t (type head hỏng ở t=0 sạch hay chỉ ở t cao?) ===
        # t=0 ~0.5  -> type head không copy nổi type sạch -> bug wiring.
        # t=0 cao, t cao ~0.5 -> không suy được loại từ vị trí bị nhiễu -> do noising vị trí.
        for ti in sorted(pred_v_by_t):
            yt = np.concatenate(true_v_by_t[ti])
            yp = np.concatenate(pred_v_by_t[ti], axis=0)
            aucs, ws = [], []
            for c in set(yt.tolist()):
                n_c = int((yt == c).sum())
                if n_c == 0 or n_c == len(yt):
                    continue
                try:
                    aucs.append(roc_auc_score(yt == c, yp[:, c]))
                    ws.append(n_c)
                except Exception:
                    pass
            t_auroc = float(np.average(aucs, weights=ws)) if aucs else float('nan')
            logger.info('[Validate per-t] t=%4d | atom auroc %.4f' % (ti, t_auroc))

        # === Đo 2: loss_pos theo từng t (thấp ở t nhỏ + cao ở t lớn = chỉ "copy", không học cấu trúc) ===
        for ti in sorted(pos_sum_by_t):
            s = pos_sum_by_t[ti]
            logger.info('[Validate per-t] t=%4d | loss_pos %.4f' % (ti, s[0] / max(s[1], 1)))

        if config.train.scheduler.type == 'plateau':
            scheduler.step(avg_loss)
        elif config.train.scheduler.type == 'warmup_plateau':
            scheduler.step_ReduceLROnPlateau(avg_loss)
        else:
            scheduler.step()

        logger.info(
            '[Validate] Iter %05d | Loss %.6f | Loss pos %.6f | Loss v %.6f e-3 | Avg atom auroc %.6f' % (
                it, avg_loss, avg_loss_pos, avg_loss_v * 1000, atom_auroc
            )
        )
        return avg_loss


    # ---- Resume: nạp lại để train tiếp (Kaggle ngắt phiên sau 12h) ----
    start_iter = 1
    best_loss, best_iter = None, None
    if args.resume is not None:
        logger.info(f'Resuming from checkpoint: {args.resume}')
        ckpt_resume = torch.load(args.resume, map_location=args.device, weights_only=False)
        model.load_state_dict(ckpt_resume['model'])
        if ckpt_resume.get('optimizer') is not None:
            optimizer.load_state_dict(ckpt_resume['optimizer'])
        if ckpt_resume.get('scheduler') is not None:
            scheduler.load_state_dict(ckpt_resume['scheduler'])
        start_iter = ckpt_resume.get('iteration', 0) + 1
        best_loss = ckpt_resume.get('best_loss', None)
        best_iter = ckpt_resume.get('best_iter', None)
        logger.info(f'  -> tiếp tục từ iter {start_iter} '
                    f'(best_loss={best_loss}, best_iter={best_iter})')

    def save_ckpt(it, fname):
        # Lưu ĐỦ state để resume: model + optimizer + scheduler + iteration + best.
        torch.save({
            'config': config,
            'model': model.state_dict(),
            'optimizer': optimizer.state_dict(),
            'scheduler': scheduler.state_dict(),
            'iteration': it,
            'best_loss': best_loss,
            'best_iter': best_iter,
        }, os.path.join(ckpt_dir, fname))

    try:
        for it in range(start_iter, config.train.max_iters + 1):
            # with torch.autograd.detect_anomaly():
            train(it)

            # Lưu last.pt định kỳ -> luôn resume được dù bị ngắt giữa 2 lần val
            if it % args.ckpt_freq == 0:
                save_ckpt(it, 'last.pt')

            if it % config.train.val_freq == 0 or it == config.train.max_iters:
                val_loss = validate(it)
                if best_loss is None or val_loss < best_loss:
                    logger.info(f'[Validate] Best val loss achieved: {val_loss:.6f}')
                    best_loss, best_iter = val_loss, it
                    save_ckpt(it, '%d.pt' % it)
                else:
                    logger.info(f'[Validate] Val loss is not improved. '
                                f'Best val loss: {best_loss:.6f} at iter {best_iter}')
                # cập nhật last.pt sau val (gồm best_loss/best_iter mới nhất)
                save_ckpt(it, 'last.pt')
    except KeyboardInterrupt:
        logger.info('Terminating...')