infer_case_unoptimized.py

import argparse
import time
from dataset import PeptidePairPicCaseDataset, encode_sequence
from network import DMutaPeptideCNN
from train import move_to_device
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
import numpy as np
from utils import set_seed
import pandas as pd

parser = argparse.ArgumentParser(description='resnet26')
# model setting
parser.add_argument('--model', type=str, default='resnet34',
                    help='resnet34 resnet50 densenet')
parser.add_argument('--q-encoder', dest='q_encoder', type=str, default='cnn',
                    help='lstm mamba mla')
parser.add_argument('--channels', type=int, default=16)
parser.add_argument("--side-enc", dest='side_enc', type=str, default='lstm',
                    help="use side features")
parser.add_argument('--fusion', type=str, default='att',
                    help='mlp att')
parser.add_argument('--glob-feat', dest='glob_feat', action='store_true', default=False,
                    help="use global features")
parser.add_argument('--non-siamese', dest='non_siamese', action='store_true', default=False,
                    help="use non-siamese architecture")

# task & dataset setting
parser.add_argument('--task', type=str, default='cls',
                    help='reg or cls')
parser.add_argument('--one-way', action='store_true', dest='one_way', default=False,
                    help='use one-way constructed dataset')
parser.add_argument('--max-length', dest='max_length', type=int, default=30,
                    help='Max length for sequence filtering')
parser.add_argument('--resize', type=int, default=[768], nargs='+',
                    help='resize the image')
parser.add_argument('--split', type=int, default=5,
                    help="Split k fold in cross validation (default: 5)")
parser.add_argument('--seed', type=int, default=1,
                    help="Seed for model initialization (default: 1)")
parser.add_argument('--pcs', action='store_true', default=False,
                    help='Consider protease cut site')
parser.add_argument('--mix-pcs', dest='mix_pcs', action='store_true', default=False,
                    help='Consider protease cut site')

# training setting
parser.add_argument('--gpu', type=int, default=0,
                    help='GPU index to use, -1 for CPU (default: 0)')
parser.add_argument('--batch-size', type=int, dest='batch_size', default=32,
                    help='input batch size for training (default: 128)')
parser.add_argument('--epochs', type=int, default=50,
                    help='number of epochs to train (default: 100)')
parser.add_argument('--lr', type=float, default=0.001,
                    help='learning rate (default: 0.001)')
parser.add_argument('--decay', type=float, default=0.0005,
                    help='weight decay (default: 0.0005)')
parser.add_argument('--pretrain', type=str, dest='pretrain', default='',
                    help='path of the pretrain model')
parser.add_argument('--metric-avg', type=str, dest='metric_avg', default='macro',
                    help='metric average type')

parser.add_argument('--loss', type=str, default='ce',
                    help='loss function')
parser.add_argument('--dir', action='store_true', default=False,
                    help='use DIR')

parser.add_argument('--simple', dest='simple', action='store_true', default=False)
parser.add_argument('--llm-data', dest='llm_data', action='store_true', default=False)

# Case Study Specific
parser.add_argument('--case', type=str, default='r2',
                    help='case to infer')
parser.add_argument('--use-ft', dest='use_ft', action='store_true', default=False)

args = parser.parse_args()

if args.llm_data:
    args.simple = True

if args.simple:
    args.one_way = True

if args.mix_pcs:
    args.pcs = 'mix'

if args.gpu != -1:
    torch.backends.cudnn.benchmark = True
    torch.set_float32_matmul_precision('high')


def load_model(args, weight_path, device):
    model = DMutaPeptideCNN(q_encoder=args.q_encoder, classes=args.classes, channels=args.channels, dir=args.dir, gf=args.glob_feat, side_enc=args.side_enc, fusion=args.fusion, non_siamese=args.non_siamese).to(device).eval()
    model.load_state_dict(torch.load(weight_path, map_location=device), strict=False)
    model.compile()
    return model


def main():
    set_seed(args.seed)
    if args.task == 'reg':
        args.classes = 1
    elif args.task == 'cls':
        args.classes = 2
    else:
        raise NotImplementedError("unimplemented task")
    weight_dir = f'./run-{args.task}/{args.q_encoder}{f"-non-siamese" if args.non_siamese else ""}-{args.fusion}-{args.channels}{f"-{args.side_enc}" if args.side_enc else ""}{"-mixpcs" if args.mix_pcs else ""}{"-pcs" if args.pcs==True else ""}{"-simple" if args.simple else ""}{"-llm" if args.llm_data else ""}{"-" + "x".join(str(n) for n in args.resize) if args.resize else ""}{"-gf" if args.glob_feat else ""}{"-oneway" if args.one_way else ""}-{args.loss + "-dir" if args.dir else args.loss}-{str(args.batch_size)}-{str(args.lr)}-{str(args.epochs)}'
    device = torch.device("cpu" if args.gpu == -1 or not torch.cuda.is_available() else f"cuda:{args.gpu}")
    print(weight_dir)
    print(device)
    
    test_set = PeptidePairPicCaseDataset(case=args.case, pad_length=args.max_length, side_enc=args.side_enc, pcs=True, resize=args.resize, gf=args.glob_feat)
    test_loader = DataLoader(test_set, batch_size=128, shuffle=False, num_workers=16, pin_memory=True)
    # test_loader = DataLoader(test_set, batch_size=192, shuffle=False, num_workers=8)
    
    models = [load_model(args, f'{weight_dir}/model_{i}{"_ft" if args.use_ft else ""}.pth', device) for i in range(args.split)]

    all_seqs = []
    logits_batches = []  # 存放每个 batch 的 [m,B,2] avg_logits (CPU 上)

    start_time = time.time()

    with torch.no_grad():
        for x, gt in test_loader:
            # x: [B, ...] on CPU pin memory，gt: tuple of B strings
            x = move_to_device(x, device, non_blocking=True)
            # x = move_to_device(x, device)

            # 1) 记录 5 个模型的 logits
            # logits: [m,B,2]
            logits = torch.zeros(len(models), len(gt), args.classes, device=device)
            for i, m in enumerate(models):
                logits[i] = m(x)

            # 3) 立刻搬到 CPU（pin_memory 下可以 non_blocking）
            logits_batches.append(logits.cpu())
            all_seqs.extend(gt)

    # 拼接成 [n,2]，n = sum(batch_size)
    all_logits = torch.cat(logits_batches, dim=1)  # [m,n,2]
    
    if args.task == 'reg':
        preds = all_logits.mean(0).squeeze().tolist()
    elif args.task == 'cls':
        # 最后一次性 softmax，取正类概率
        preds = torch.softmax(all_logits, dim=-1)[:, :, 1].mean(0).tolist()

    consumed_time = time.time() - start_time
    print(f'total consumed time: {consumed_time} s')
    print(f'time per sample: {consumed_time / len(test_set)} s')

    # 保存到 CSV
    df = pd.DataFrame({
        "seq":  all_seqs,
        "pred": preds,
    })

    df.to_csv(f'{weight_dir}/preds_case.csv', index=False)




if __name__ == '__main__':
    main()