uda.py

import argparse
import json
import logging
import os
import time
import math
from copy import deepcopy

from dataset import PeptidePairCaseDataset, PeptidePairDataset, PeptidePairPicCaseDataset, PeptidePairPicDataset
from network import DMutaPeptide, DMutaPeptideCNN
from train import move_to_device, update_ce_loss_weight
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms.v2 as T
from torch.utils.data import DataLoader, Subset, RandomSampler
import numpy as np
from loss import MLCE, SuperLoss, LogCoshLoss, BMCLoss
from utils import set_seed, zip_restart_dataloader as zrd
from random import shuffle
from torchmetrics import MeanAbsoluteError, RelativeSquaredError, PearsonCorrCoef, KendallRankCorrCoef, Accuracy, F1Score, AveragePrecision, AUROC


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='lstm',
                    help='lstm mamba mla')
parser.add_argument("--side-enc", dest='side_enc', type=str, default=None,
                    help="use side features")
parser.add_argument('--channels', type=int, default=256)
parser.add_argument('--fusion', type=str, default='att',
                    help='mlp att diff')
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='reg',
                    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('--split', type=int, default=5,
                    help="Split k fold in cross validation (default: 5)")
parser.add_argument('--seed', type=int, default=42,
                    help="Seed (default: 1)")
parser.add_argument('--pcs', action='store_true', default=False,
                    help='Consider protease cleavage site')
parser.add_argument('--mix-pcs', dest='mix_pcs', action='store_true', default=False,
                    help='Consider protease cleavage site')
parser.add_argument('--resize', type=int, default=[768], nargs='+',
                    help='resize the image')
# parser.add_argument('--llm-data', action='store_true', default=False,
#                     help='Use LLM augmentation data')

# 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='mse',
                    help='loss function')
parser.add_argument('--dir', action='store_true', default=False,
                    help='use DIR')

parser.add_argument('--case', type=str, default='r2')
parser.add_argument('--simple', action='store_true', default=False)
parser.add_argument('--pt', action='store_true', default=False)
parser.add_argument('--uda-epochs', '-ue', dest='uda_epochs', default=50)

args = parser.parse_args()

if args.simple:
    args.one_way = True

class GaussianNoise(nn.Module):
    def __init__(self, mean=0., sigma=0.15):
        super(GaussianNoise, self).__init__()
        self.mean = mean
        self.sigma = sigma

    def forward(self, x):
        return x + torch.randn_like(x) * self.sigma + self.mean

# --------------------
# 强/弱增强
# --------------------
strong_transforms = T.Compose([
    T.RandomResizedCrop(args.resize, scale=(0.7, 1.0)),
    T.RandomGrayscale(0.2),
    GaussianNoise(0., 0.4),
])
weak_transforms = T.Compose([
    T.RandomResizedCrop(args.resize, scale=(0.9, 1.0)),
    GaussianNoise(0., 0.05),
])

def strong_aug(x, device=torch.device('cpu')):
    return aug_and_move(x, strong_transforms, 0.2, device, True)

def weak_aug(x, device=torch.device('cpu')):
    return aug_and_move(x, weak_transforms, 0.05, device, True)

def aug_and_move(x, transforms: T.Transform, seq_noise=0.05, device=torch.device('cpu'), non_blocking=False):
    if isinstance(x, (tuple, list)):
        return type(x)(aug_and_move(x_i, transforms, seq_noise, device, non_blocking) for x_i in x)
    if len(x.shape) == 3:
        return (x + torch.randn_like(x) * seq_noise).to(device, non_blocking=non_blocking)
    else:
        # return transforms(x.to(device, non_blocking=non_blocking))
        return torch.stack([transforms(s) for s in x.to(device, non_blocking=non_blocking)], dim=0)

# --------------------
#  EMA 更新函数
# --------------------
def update_ema(student_model: nn.Module, teacher_model: nn.Module, alpha):
    for s_param, t_param in zip(student_model.parameters(), teacher_model.parameters()):
        t_param.data = t_param.data.mul_(alpha).add_(s_param.data, alpha=(1 - alpha))

# --------------------
#  学习率与一致性权重 ramp-up
# --------------------
def sigmoid_rampup(current, rampup_length):
    if rampup_length == 0:
        return 1.0
    else:
        current = max(0.0, min(1.0, current / rampup_length))
        return float(math.exp(-5.0 * (1.0 - current) * (1.0 - current)))
    
# --------------------
#  交叉熵一致性损失函数
# --------------------
def consistency_loss_ce(s_pred, t_pred, threshold=None):
    """
    s_pred: 学生模型 logits (B, C)
    t_pred: 教师模型 logits (B, C)
    mask: 可选的布尔 mask，True 表示该样本参与 loss 计算
    """
    # 1) 计算 teacher softmax 概率
    probs = F.softmax(t_pred.detach(), dim=1)  # (B, C)
    # 2) 生成伪标签
    max_probs, pseudo_labels = probs.max(dim=1)  # (B,), (B,)
    # 3) 可选：基于置信度做 mask
    if threshold is None:
        mask = torch.ones_like(max_probs, dtype=torch.float)
    else:
        mask = max_probs.ge(threshold).float()  # (B,) 0/1
    # 4) 计算交叉熵，并按 mask 加权
    loss = F.cross_entropy(s_pred, pseudo_labels, reduction='none')  # (B,)
    loss = (loss * mask).sum() / (mask.sum().clamp(min=1.0))
    return loss


def main():
    set_seed(args.seed)
    if args.task == 'reg':
        args.classes = 1
        if args.loss == "mse" or args.loss in ['ce']:
            args.loss = 'mse'
            criterion = nn.MSELoss()
            criterion_cons = criterion
        else:
            raise NotImplementedError("unimplemented regression task loss function")
    elif args.task == 'cls':
        args.classes = 2
        if args.loss == 'ce' or args.loss in ['mse', 'smoothl1', 'super']:
            args.loss = 'ce'
            criterion = nn.CrossEntropyLoss()
            criterion_cons = consistency_loss_ce
            # criterion_cons = nn.MSELoss()
        else:
            raise NotImplementedError("unimplemented classification task loss function")
    else:
        raise NotImplementedError("unimplemented task")

    if args.q_encoder in ['cnn', 'rn18']:
        weight_dir = f'./run-{args.task}/{args.q_encoder}{"-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 ""}{"-" + "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)}/uda_{args.case}'
    else:
        weight_dir = f'./run-{args.task}/{args.q_encoder}{f"-non-siamese" if args.non_siamese else ""}-{args.fusion}-{args.channels}{"-simple" if args.simple 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)}/uda_{args.case}'

    if not os.path.exists(weight_dir):
        os.makedirs(weight_dir)
    
    logging.basicConfig(handlers=[
        logging.FileHandler(filename=os.path.join(weight_dir, "training.log"), encoding='utf-8', mode='w+'),
        logging.StreamHandler()],
        format="%(asctime)s: %(message)s", datefmt="%F %T", level=logging.INFO)

    logging.info(f'saving_dir: {weight_dir}')
    
    with open(os.path.join(weight_dir, "config.json"), "w") as f:
        f.write(json.dumps(vars(args)))

    device = torch.device("cpu" if args.gpu == -1 or not torch.cuda.is_available() else f"cuda:{args.gpu}")

    logging.info('Loading Training Dataset')
    if args.q_encoder in ['cnn', 'rn18']:
        unlabel_set = PeptidePairPicCaseDataset(case=args.case, pad_length=args.max_length, side_enc=args.side_enc, pcs=args.pcs, resize=args.resize, gf=args.glob_feat)
    else:
        unlabel_set = PeptidePairCaseDataset(case=args.case, pad_length=args.max_length, gf=args.glob_feat)
    unlabel_loader = DataLoader(unlabel_set, batch_size=args.batch_size // 2, shuffle=True, drop_last=True, num_workers=16,  pin_memory=True)

    if args.q_encoder in ['cnn', 'rn18']:
        label_set = PeptidePairPicDataset(mode='train', pad_length=args.max_length, task=args.task, side_enc=args.side_enc, pcs=args.pcs, resize=args.resize, one_way=args.one_way, gf=args.glob_feat)
    else:
        label_set = PeptidePairDataset(mode='train', pad_length=args.max_length, task=args.task, one_way=args.one_way, gf=args.glob_feat)
    label_loader = DataLoader(label_set, batch_size=args.batch_size // 2, shuffle=True, drop_last=True, num_workers=16, pin_memory=True)

    if args.case == 'r2':
        logging.info('Loading Validation Dataset')
        if args.q_encoder in ['cnn', 'rn18']:
            val_set = PeptidePairPicDataset(mode='r2_case', pad_length=args.max_length, task=args.task, side_enc=args.side_enc, pcs=args.pcs, resize=args.resize, gf=args.glob_feat)
        else:
            val_set = PeptidePairDataset(mode='r2_case', pad_length=args.max_length, task=args.task, gf=args.glob_feat)
        val_loader = DataLoader(val_set, batch_size=args.batch_size * 2, shuffle=False, num_workers=16, pin_memory=True)
        metric_funcs = get_metric_funcs(args.task, device)
    else:
        val_loader = None

    best_val_metric = -float('inf')
    logging.info(f"Start UDA training")
    weights_path = f"{weight_dir}/model_uda_{'{role}'}.pth"

    if args.q_encoder in ['cnn', 'rn18']:
        student = 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).train()
    else:
        student = DMutaPeptide(q_encoder=args.q_encoder, classes=args.classes, channels=args.channels, dir=args.dir, gf=args.glob_feat, fusion=args.fusion, non_siamese=args.non_siamese).to(device).train()
    if args.pt:
        student.load_state_dict(torch.load(os.path.join(os.path.dirname(weight_dir), 'model_0_test.pth'), map_location=device))
    teacher = deepcopy(student).to(device).eval()
    for param in teacher.parameters():
        param.requires_grad = False
    
    optimizer = torch.optim.AdamW(student.parameters(), lr=1e-4)
    
    global_step = 0
    rampup_length = 1500
    for epoch in range(1, args.uda_epochs+1):
        train_loss = []
        for (x_l, y_l), (x_u, _) in zrd(label_loader, unlabel_loader):
            x_l = move_to_device(x_l, device, non_blocking=True)
            y_l = y_l.to(device, non_blocking=True)

            # -- 1) 有监督损失 -- 
            pred_l = student(x_l)
            if args.loss == 'ce':
                update_ce_loss_weight(criterion, y_l, num_classes=2, device=device)
            loss_sup = criterion(pred_l, y_l)

            # -- 2) 一致性损失 -- 
            with torch.no_grad():
                t_pred = teacher(weak_aug(x_u, device))

            s_pred = student(strong_aug(x_u, device))

            loss_cons = criterion_cons(s_pred, t_pred)
            
            # -- 3) 总损失 -- 
            λ = 1.0 * sigmoid_rampup(global_step, rampup_length)
            loss = loss_sup + λ * loss_cons

            # -- 4) 学生模型更新 --
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            # -- 5) 教师 EMA 更新 -- 
            alpha = 0.99
            update_ema(student, teacher, alpha)

            global_step += 1
            train_loss.append(loss.item())
        
        train_loss = sum(train_loss) / len(train_loss)

        if val_loader:
            with torch.no_grad():
                val_pred, val_gt = [], []
                for x, gt in val_loader:
                    x = move_to_device(x, device, non_blocking=True)
                    out = teacher(x)
                    # out = student(x)
                    val_pred.append(out)
                    val_gt.append(gt.to(device, non_blocking=True))
                val_pred = torch.cat(val_pred, dim=0)
                val_gt = torch.cat(val_gt, dim=0)

            if args.task == 'cls':
                val_pred = torch.softmax(val_pred, dim=1)[:, 1]
                val_ap = metric_funcs['ap'](val_pred, val_gt).item()
                val_auc = metric_funcs['auc'](val_pred, val_gt).item()
                val_f1 = metric_funcs['f1'](val_pred, val_gt).item()
                val_acc = metric_funcs['acc'](val_pred, val_gt).item()
                val_metric = val_ap + val_auc
                logging.info(f'Epoch {epoch} Train Loss: {train_loss:.4f} Val: ap: {val_ap:.4f} auc: {val_auc:.4f} f1: {val_f1:.4f} acc: {val_acc:.4f}')
            elif args.task == 'reg':
                val_mae = metric_funcs['mae'](val_pred, val_gt).item()
                val_rse = metric_funcs['rse'](val_pred, val_gt).item()
                val_pcc = metric_funcs['pcc'](val_pred, val_gt).item()
                val_kcc = metric_funcs['kcc'](val_pred, val_gt).item()
                val_metric = val_pcc + val_kcc - val_mae - val_rse
                logging.info(f'Epoch {epoch} Train Loss: {train_loss:.4f} Val: mae: {val_mae:.4f} rse: {val_rse:.4f} pcc: {val_pcc:.4f} kcc: {val_kcc:.4f}')
            else:
                raise NotImplementedError

            if val_metric > best_val_metric:
                best_val_metric = val_metric
                logging.info(f'Epoch: {epoch} New best VAL metrics')
                torch.save(student.state_dict(), weights_path.format(role='student'))
                torch.save(teacher.state_dict(), weights_path.format(role='teacher'))
        else:
            logging.info(f'Epoch {epoch} Train Loss: {train_loss:.4}')
            if (args.task == 'reg' and train_loss > 0.199) or (args.task == 'cls' and train_loss > 0.259):
                val_metric = -train_loss
                if val_metric > best_val_metric:
                    best_val_metric = val_metric
                    torch.save(student.state_dict(), weights_path.format(role='student'))
                    torch.save(teacher.state_dict(), weights_path.format(role='teacher'))
            else:
                break
                
    logging.info('UDA training finished')
    torch.save(student.state_dict(), weights_path.format(role='student_last'))
    torch.save(teacher.state_dict(), weights_path.format(role='teacher_last'))


def get_metric_funcs(task, device):
    if task == 'reg':
        metric_funcs = {
            'mae': MeanAbsoluteError().to(device),
            'rse': RelativeSquaredError().to(device),
            'pcc': PearsonCorrCoef().to(device),
            'kcc': KendallRankCorrCoef().to(device)
        }
    elif task == 'cls':
        metric_funcs = {
            'ap': AveragePrecision(task='binary').to(device),
            'auc': AUROC(task='binary').to(device),
            'f1': F1Score(task='binary').to(device),
            'acc': Accuracy(task='binary').to(device)
        }
    else:
        raise NotImplementedError(f'Task {task} not supported')
    return metric_funcs

if __name__ == '__main__':
    main()