text_classification/train.py

# 简化版本 - 使用标准的Hugging Face方法
import torch
import pandas as pd
from torch.utils.data import random_split
from torch.utils.data import Dataset
from torch.utils.data import DataLoader
from transformers import AutoTokenizer, AutoModelForSequenceClassification, get_linear_schedule_with_warmup
from torch.optim.lr_scheduler import ReduceLROnPlateau
from tqdm import tqdm
import os
from torch.utils.data import WeightedRandomSampler
import numpy as np
from collections import Counter
import time
import gc

class MyDataset(Dataset):
    def __init__(self, file_path):
        self.file_path = file_path
        self.data = pd.read_csv(file_path)
        self.data = self.data.dropna()
        # 打印数据集基本信息
        print(f"\n数据集信息:")
        print(f"  总样本数: {len(self.data)}")
        print(f"  列名: {self.data.columns.tolist()}")

        # 统计文本长度分布
        if 'record' in self.data.columns:
            text_lengths = self.data['record'].astype(str).apply(len)
            print(f"  文本长度统计:")
            print(f"    最小: {text_lengths.min()}")
            print(f"    最大: {text_lengths.max()}")
            print(f"    平均: {text_lengths.mean():.1f}")
            print(f"    中位数: {text_lengths.median():.1f}")
            print(f"    95%分位数: {text_lengths.quantile(0.95):.1f}")

        # 统计标签分布
        if 'label' in self.data.columns:
            label_counts = Counter(self.data['label'])
            print(f"  标签分布:")
            for label, count in sorted(label_counts.items()):
                print(f"    类别 {label}: {count} ({count / len(self.data) * 100:.1f}%)")

    def __getitem__(self, item):
        return self.data.iloc[item]['record'], self.data.iloc[item]['label']

    def __len__(self):
        return len(self.data)


class EarlyStopping:
    """Early stopping机制"""

    def __init__(self, patience:int=3, min_delta:float=0, mode:str='min'):
        self.patience = patience
        self.min_delta = min_delta
        self.mode = mode
        self.counter = 0
        self.best_score = None
        self.early_stop = False

    def __call__(self, score):
        if self.best_score is None:
            self.best_score = score
            return False

        if self.mode == 'min':
            if score < self.best_score - self.min_delta:
                self.best_score = score
                self.counter = 0
            else:
                self.counter += 1
        else:  # mode == 'max'
            if score > self.best_score + self.min_delta:
                self.best_score = score
                self.counter = 0
            else:
                self.counter += 1

        if self.counter >= self.patience:
            self.early_stop = True
            return True
        return False


class SimpleTrainer:
    def __init__(self, train_csv_file,
                 batch_size:int=32,
                 val_csv_file:str=None,
                 is_english:bool=False,
                 num_classes:int=2,
                 max_length:int=256):
        # 使用标准的BERT模型
        if is_english:
            self.model_name = "bert-base-uncased"
        else:
            self.model_name = "bert-base-chinese"
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        print(f"使用设备: {self.device}")
        print(f"使用模型: {self.model_name}")

        self.batch_size = batch_size
        self.max_length = max_length
        self.__global_step = 0
        self.best_val_acc = 0.0
        self.best_val_loss = float('inf')

        # 创建checkpoint目录
        self.checkpoint_root_path = './checkpoints'
        if not os.path.exists(self.checkpoint_root_path):
            os.makedirs(self.checkpoint_root_path)

        self.best_acc_model_path = os.path.join(self.checkpoint_root_path , f'{self.model_name}_best_acc.pth')
        self.best_loss_model_path = os.path.join(self.checkpoint_root_path , f'{self.model_name}_best_loss.pth')
        # 断点续训
        self.latest_checkpoint_path = os.path.join(self.checkpoint_root_path , f'{self.model_name}_latest_checkpoint.pth')

        # 创建模型和分词器
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
        self.model = AutoModelForSequenceClassification.from_pretrained(
            self.model_name,
            num_labels=num_classes
        ).to(self.device)

        # 打印模型参数量
        total_params = sum(p.numel() for p in self.model.parameters())
        trainable_params = sum(p.numel() for p in self.model.parameters() if p.requires_grad)
        print(f"模型总参数量: {total_params:,}")
        print(f"可训练参数量: {trainable_params:,}")

        # 创建数据集
        print("\n加载数据集...")
        self.train_dataset = MyDataset(train_csv_file)
        # 计算样本权重（处理类别不平衡）
        labels = self.train_dataset.data['label'].values
        class_counts = np.bincount(labels)
        sample_weights = 1.0 / class_counts[labels]  # 实际上给每个样本定义了采样的权重
        sampler = WeightedRandomSampler(
            weights=sample_weights,
            num_samples=len(sample_weights),
            replacement=True
        )
        if val_csv_file:
            self.val_dataset = MyDataset(val_csv_file)
        else:
            # 如果没有独立验证集，需要分割
            self.train_dataset, self.val_dataset = random_split(self.train_dataset, [0.8, 0.2])
            # 从完整权重中提取训练集对应的权重
            train_indices = self.train_dataset.indices
            train_sample_weights = sample_weights[train_indices]
            sampler = WeightedRandomSampler(
                weights=train_sample_weights,
                num_samples=len(train_sample_weights),
                replacement=True
            )
        # 创建数据加载器
        self.train_loader = DataLoader(self.train_dataset,
                                       batch_size=self.batch_size,
                                       sampler=sampler,
                                       collate_fn=self.collate_func)
        self.val_loader = DataLoader(self.val_dataset,
                                     batch_size=self.batch_size,
                                     shuffle=False,
                                     collate_fn=self.collate_func)

        # 创建优化器和调度器
        self.optimizer = torch.optim.AdamW(self.model.parameters(), lr=3e-5, weight_decay=0.01)
        self.scheduler = None
        self.warmup_ratio = 0.1  # warmup步数占总步数的比例

        # Early stopping
        self.early_stopping = EarlyStopping(patience=5, min_delta=0.001, mode='min')
        # 梯度裁剪阈值
        self.max_grad_norm = 1.0

    def collate_func(self, batch_data):
        texts, labels = [], []
        for item in batch_data:
            texts.append(item[0])
            labels.append(item[1])
        
        inputs = self.tokenizer(
            texts, 
            max_length=self.max_length,
            padding='max_length', 
            truncation=True, 
            return_tensors='pt'
        )
        inputs['labels'] = torch.tensor(labels, dtype=torch.long)
        return inputs

    def train_step(self):
        self.model.train()
        total_loss = 0.0
        correct_predictions = 0.0
        total_samples = 0.0
        
        for batch_data in tqdm(self.train_loader, desc="Training"):
            batch_data = {k: v.to(self.device) for k, v in batch_data.items()}
            
            self.optimizer.zero_grad()
            outputs = self.model(**batch_data)
            loss = outputs.loss
            loss.backward()
            # 添加梯度裁剪
            torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.max_grad_norm)

            self.optimizer.step()
            self.scheduler.step()  # 更新学习率
            
            total_loss += loss.item()
            predictions = torch.argmax(outputs.logits, dim=-1)
            correct_predictions += (predictions == batch_data['labels']).float().sum().item()
            total_samples += len(batch_data['labels'])
            
        avg_loss = total_loss / len(self.train_loader)
        accuracy = correct_predictions / total_samples
        return avg_loss, accuracy

    def val_step(self):
        self.model.eval()
        total_loss = 0.0
        correct_predictions = 0.0
        total_samples = 0.0
        
        with torch.no_grad():
            for batch_data in tqdm(self.val_loader, desc="Validating"):
                batch_data = {k: v.to(self.device) for k, v in batch_data.items()}
                
                outputs = self.model(**batch_data)
                total_loss += outputs.loss.item()
                
                predictions = torch.argmax(outputs.logits, dim=-1)
                correct_predictions += (predictions == batch_data['labels']).float().sum().item()
                total_samples += len(batch_data['labels'])
                
        avg_loss = total_loss / len(self.val_loader)
        accuracy = correct_predictions / total_samples
        return avg_loss, accuracy

    def train_and_validate(self, num_epoch):
        # 创建学习率调度器(带warmup的线性衰减)
        total_steps = len(self.train_loader) * num_epoch
        warmup_steps = int(total_steps * self.warmup_ratio)

        self.scheduler = get_linear_schedule_with_warmup(
            self.optimizer,
            num_warmup_steps=warmup_steps,
            num_training_steps=total_steps
        )

        print(f"\n训练配置:")
        print(f"  总训练步数: {total_steps}")
        print(f"  Warmup步数: {warmup_steps}")
        print(f"  初始学习率: {self.optimizer.param_groups[0]['lr']}")
        print(f"  Batch size: {self.batch_size}")
        print(f"  Max length: {self.max_length}")
        print(f"  梯度裁剪: {self.max_grad_norm}")

        print("开始训练...")
        for epoch in range(num_epoch):
            print(f"\n{'='*60}")
            print(f"Epoch {epoch+1}/{num_epoch}")
            print(f"{'='*60}")

            train_loss, train_acc = self.train_step()
            val_loss, val_acc = self.val_step()

            # 获取当前学习率
            current_lr = self.optimizer.param_groups[0]['lr']

            print(f'Epoch {epoch+1}/{num_epoch}:')
            print(f'  Train Loss: {train_loss:.4f}, Acc: {train_acc:.4f}')
            print(f'  Val Loss: {val_loss:.4f}, Acc: {val_acc:.4f}')
            print(f'  Learning Rate: {current_lr:.2e}')

            # 保存最佳模型 (基于验证准确率)
            if val_acc > self.best_val_acc:
                self.best_val_acc = val_acc
                torch.save(self.model.state_dict(), self.best_acc_model_path)
                print(f"  ✓ 保存了新的最佳准确率模型，验证准确率: {self.best_val_acc:.4f}")

            # 保存最低验证损失模型
            if val_loss < self.best_val_loss:
                self.best_val_loss = val_loss
                torch.save(self.model.state_dict(), self.best_loss_model_path)
                print(f"  ✓ 保存了新的最低损失模型，验证损失: {self.best_val_loss:.4f}")

            # Early stopping检查
            if self.early_stopping(val_loss):
                print(f"\n⚠ Early stopping触发! 在epoch {epoch+1}停止训练")
                print(f"  验证损失已连续{self.early_stopping.patience}个epoch没有改善")
                break

        print(f"\n{'='*60}")
        print("训练完成!")
        print(f"{'='*60}")
        print(f"最佳验证准确率: {self.best_val_acc:.4f}")
        print(f"最低验证损失: {self.best_val_loss:.4f}")
        print(f"模型保存路径:")
        print(f"  最佳准确率模型: {self.best_acc_model_path}")
        print(f"  最低损失模型: {self.best_loss_model_path}")

if __name__ == '__main__':
    trainer_en = SimpleTrainer(train_csv_file='data_en.csv',
                            batch_size=32,
                            is_english=True,
                            max_length=256)  # 增加max_length
    trainer_en.train_and_validate(num_epoch=20)

    del trainer_en
    gc.collect()  # 强制 Python 垃圾回收
    torch.cuda.empty_cache()  # 清空 CUDA 缓存
    time.sleep(2)  # 短暂等待即可，不需要太久

    trainer = SimpleTrainer(train_csv_file='data_zh.csv',
                            batch_size=32,
                            is_english=False,
                            max_length=256)  # 增加max_length
    trainer.train_and_validate(num_epoch=20)