train.py

import os

os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"

import numpy as np
import torch
import torch.nn as nn
from torch import optim
from torch.utils.data import DataLoader
from torch.optim.lr_scheduler import _LRScheduler
from torchvision import transforms
from scipy.stats import pearsonr, spearmanr
from sklearn.metrics import r2_score
from transformers import BertTokenizer

from dataset import ic_dataset_four_branch_fusion
from options import args
from model.TSGNet import TSGNet


# ==========================================
# Utilities (原 utils.py 内容)
# ==========================================
class WarmUpLR(_LRScheduler):
    def __init__(self, optimizer, total_iters, last_epoch=-1):
        self.total_iters = total_iters
        super().__init__(optimizer, last_epoch)

    def get_lr(self):
        return [base_lr * self.last_epoch / (self.total_iters + 1e-8) for base_lr in self.base_lrs]


def evaInfo(score, label):
    score = np.array(score)
    label = np.array(label)
    RMAE = np.sqrt(np.abs(score - label).mean())
    RMSE = np.sqrt(np.mean(np.abs(score - label) ** 2))
    Pearson = np.float64(pearsonr(label, score)[0])
    Spearmanr = np.float64(spearmanr(label, score)[0])
    tx = score
    ty = label
    ccc = (2 * Pearson * tx.std() * ty.std()) / (np.power(tx.mean() - ty.mean(), 2) + tx.var() + ty.var())
    R2 = r2_score(label, score)
    info = (
        f"\n[Eval Results] "
        f"RMSE: {RMSE:.4f} | RMAE: {RMAE:.4f} | "
        f"Pearson: {Pearson:.4f} | Spearman: {Spearmanr:.4f} | "
        f"CCC: {ccc:.4f} | R2: {R2:.4f}"
    )
    infoList = (RMSE, RMAE, Pearson, Spearmanr, ccc)
    return info, infoList


# ==========================================
# Trainer 类重构
# ==========================================
class Trainer:
    def __init__(self, args):
        self.args = args
        self.device = torch.device("cuda:{}".format(args.gpu_id))

        self._build_dataloaders()
        self._build_model()
        self._build_optimizers()

        self.bestSum = 100000

    def _build_dataloaders(self):
        """构建训练和测试数据集及加载器"""
        trainTransform = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])
        testTransform = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])

        tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')

        trainDataset = ic_dataset_four_branch_fusion(
            txt_path="./IC9600/train_merged.txt",
            jsonl_path="./IC9600/api_results.jsonl",  # 请根据实际路径修改
            img_path="./IC9600/images/",
            tokenizer=tokenizer,
            transform=trainTransform
        )
        self.trainDataLoader = DataLoader(
            trainDataset,
            batch_size=self.args.batch_size,
            num_workers=self.args.num_workers,
            shuffle=True
        )

        testDataset = ic_dataset_four_branch_fusion(
            txt_path="./IC9600/test_merged.txt",
            jsonl_path="./IC9600/api_results.jsonl",  # 请根据实际路径修改
            img_path="./IC9600/images/",
            tokenizer=tokenizer,
            transform=testTransform
        )
        self.testDataLoader = DataLoader(
            testDataset,
            batch_size=self.args.batch_size,
            num_workers=self.args.num_workers,
            shuffle=False
        )

    def _build_model(self):
        """初始化模型"""
        if not os.path.exists(self.args.ck_save_dir):
            os.mkdir(self.args.ck_save_dir)

        self.model = TSGNet()
        self.model.to(self.device)

        for name, param in self.model.named_parameters():
            if param.requires_grad:
                print(name)

    def _build_optimizers(self):
        """初始化损失函数、优化器和学习率调度器"""
        self.loss_function = nn.MSELoss()

        self.optimizer = optim.AdamW(
            filter(lambda p: p.requires_grad, self.model.parameters()),
            lr=self.args.lr,
            betas=(0.9, 0.999),
            eps=1e-08,
            weight_decay=0.01,
            amsgrad=False
        )

        self.scheduler = optim.lr_scheduler.ReduceLROnPlateau(
            self.optimizer,
            mode='min',
            factor=0.5,
            patience=8,
            threshold=1e-4,
            threshold_mode='rel',
            cooldown=0,
            min_lr=0,
            eps=1e-8
        )

        iter_per_epoch = len(self.trainDataLoader)
        if self.args.warm > 0:
            self.warmup_scheduler = WarmUpLR(self.optimizer, iter_per_epoch * self.args.warm)
        else:
            self.warmup_scheduler = None

    def _prepare_data(self, batch):
        """将 batch 数据移动到设备上"""
        data = {k: v.to(self.device) for k, v in batch.items()}
        return data

    def train_epoch(self, epoch):
        """执行单个 epoch 的训练"""
        self.model.train()
        for batch_index, batch in enumerate(self.trainDataLoader):
            data = self._prepare_data(batch)

            self.optimizer.zero_grad()
            score = self.model(
                data['img'],
                data['g_ids'], data['g_mask'],
                data['n_ids'], data['n_mask'],
                data['s_ids'], data['s_mask'],
                data['a_ids'], data['a_mask']
            )
            loss = self.loss_function(score, data['label'])
            loss.backward()
            self.optimizer.step()

            if self.warmup_scheduler and epoch <= self.args.warm:
                self.warmup_scheduler.step()

            if (batch_index + 1) % (len(self.trainDataLoader) // 3) == 0:
                print('Training Epoch: {epoch} [{trained_samples}/{total_samples}]\tloss: {:0.4f}\tLR: {:0.6f}'.format(
                    loss.item(),
                    self.optimizer.param_groups[0]['lr'],
                    epoch=epoch,
                    trained_samples=batch_index * self.args.batch_size + len(data['img']),
                    total_samples=len(self.trainDataLoader.dataset)
                ))

    def evaluate(self):
        """在测试集上进行评估"""
        self.model.eval()
        all_scores = []
        all_labels = []

        for batch in self.testDataLoader:
            data = self._prepare_data(batch)
            with torch.no_grad():
                score = self.model(
                    data['img'],
                    data['g_ids'], data['g_mask'],
                    data['n_ids'], data['n_mask'],
                    data['s_ids'], data['s_mask'],
                    data['a_ids'], data['a_mask']
                )
            all_scores += score.tolist()
            all_labels += data['label'].tolist()

        info, infoList = evaInfo(score=all_scores, label=all_labels)
        print(info)
        return infoList

    def run(self):
        """启动整个训练流程"""
        for epoch in range(1, self.args.epoch + 1):
            self.train_epoch(epoch)

            current_sum = 0
            infoList = self.evaluate()
            current_sum += infoList[0]
            current_sum += infoList[1]
            current_sum += 1 - np.abs(infoList[2])
            current_sum += 1 - np.abs(infoList[3])
            current_sum += 1 - np.abs(infoList[4])

            if current_sum < self.bestSum:
                torch.save(self.model.state_dict(), os.path.join(self.args.ck_save_dir, 'ck_best'))
                self.bestSum = current_sum

            if epoch > self.args.warm:
                self.scheduler.step(current_sum)

            torch.save(self.model.state_dict(), os.path.join(self.args.ck_save_dir, 'ck_TSGNet{}.pth'.format(epoch)))


if __name__ == "__main__":
    trainer = Trainer(args)
    trainer.run()