825 lines
27 KiB
Python
825 lines
27 KiB
Python
"""
|
||
RoBERTa-wwm-ext-large 商品分类微调脚本 (ModelScope版)
|
||
硬件配置: 8核CPU / 32GB内存 / 24GB显存
|
||
模型来源: ModelScope
|
||
"""
|
||
|
||
import pandas as pd
|
||
import numpy as np
|
||
from sklearn.model_selection import train_test_split
|
||
from sklearn.preprocessing import LabelEncoder
|
||
from sklearn.utils.class_weight import compute_class_weight
|
||
from sklearn.metrics import accuracy_score, f1_score, classification_report
|
||
from transformers import (
|
||
AutoTokenizer,
|
||
AutoModelForSequenceClassification,
|
||
Trainer,
|
||
TrainingArguments,
|
||
EarlyStoppingCallback,
|
||
set_seed,
|
||
TrainerCallback
|
||
)
|
||
from transformers.modeling_outputs import SequenceClassifierOutput
|
||
import torch
|
||
import torch.nn.functional as F
|
||
from torch.utils.data import Dataset
|
||
from tqdm import tqdm
|
||
import warnings
|
||
import re
|
||
import joblib
|
||
import os
|
||
import time
|
||
import glob
|
||
from datetime import datetime
|
||
from torch.utils.tensorboard import SummaryWriter
|
||
from safetensors.torch import save_file as safe_save_file
|
||
|
||
# ==================== ModelScope 导入 ====================
|
||
try:
|
||
from modelscope.hub.snapshot_download import snapshot_download
|
||
MODELSCOPE_AVAILABLE = True
|
||
print("✅ ModelScope 已安装")
|
||
except ImportError:
|
||
MODELSCOPE_AVAILABLE = False
|
||
print("⚠️ ModelScope 未安装，将使用HuggingFace缓存")
|
||
print(" 安装命令: pip install modelscope")
|
||
|
||
|
||
# ==================== 1. 参数配置 ====================
|
||
class Config:
|
||
# ===== 模型配置 =====
|
||
MODEL_ID = "dienstag/chinese-roberta-wwm-ext"
|
||
|
||
# 本地缓存目录
|
||
CACHE_DIR = "./modelscope_cache"
|
||
|
||
# 是否使用ModelScope下载
|
||
USE_MODELSCOPE = True
|
||
|
||
# ===== 训练参数 =====
|
||
MAX_LENGTH = 64
|
||
BATCH_SIZE = 32
|
||
GRADIENT_ACCUMULATION_STEPS = 2
|
||
NUM_EPOCHS = 5
|
||
LEARNING_RATE = 2e-5
|
||
WARMUP_RATIO = 0.1
|
||
WEIGHT_DECAY = 0.01
|
||
FP16 = True
|
||
EARLY_STOPPING_PATIENCE = 3
|
||
SAVE_TOTAL_LIMIT = 5
|
||
|
||
# 日志和保存步数
|
||
LOGGING_STEPS = 50
|
||
EVAL_SAVE_STEPS = 200
|
||
|
||
# 数据加载
|
||
NUM_WORKERS = 6
|
||
TEST_SIZE = 0.2
|
||
|
||
# 路径
|
||
DATA_FILE = "train_v4.xlsx"
|
||
OUTPUT_DIR = "./results/roberta_large"
|
||
LOG_DIR = "./logs/roberta_large"
|
||
SAVE_DIR = "./saved_model/roberta_large"
|
||
|
||
# 随机种子
|
||
SEED = 42
|
||
|
||
# 设备
|
||
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
|
||
|
||
# 评估指标配置
|
||
METRIC_FOR_BEST_MODEL = "f1_macro"
|
||
GREATER_IS_BETTER = True
|
||
|
||
|
||
# ==================== 2. ModelScope 模型加载 ====================
|
||
def download_model_from_modelscope(model_id, cache_dir):
|
||
"""从ModelScope下载模型"""
|
||
if not MODELSCOPE_AVAILABLE:
|
||
raise ImportError("请先安装 modelscope: pip install modelscope")
|
||
|
||
print(f"📥 从ModelScope下载模型: {model_id}")
|
||
print(f"📁 缓存目录: {cache_dir}")
|
||
|
||
try:
|
||
model_dir = snapshot_download(
|
||
model_id,
|
||
cache_dir=cache_dir,
|
||
revision='master'
|
||
)
|
||
print(f"✅ 模型下载完成: {model_dir}")
|
||
return model_dir
|
||
except Exception as e:
|
||
print(f"⚠️ ModelScope下载失败: {e}")
|
||
print(" 尝试使用HuggingFace缓存...")
|
||
return None
|
||
|
||
|
||
def get_model_path(model_id, cache_dir, use_modelscope=True):
|
||
"""获取模型路径"""
|
||
if use_modelscope and MODELSCOPE_AVAILABLE:
|
||
try:
|
||
local_dir = download_model_from_modelscope(model_id, cache_dir)
|
||
if local_dir and os.path.exists(local_dir):
|
||
return local_dir
|
||
except Exception as e:
|
||
print(f"⚠️ ModelScope加载失败: {e}")
|
||
|
||
print(f"🔄 使用HuggingFace加载: {model_id}")
|
||
return model_id
|
||
|
||
|
||
# ==================== 3. 工具函数 ====================
|
||
def print_gpu_memory():
|
||
"""打印GPU显存使用情况"""
|
||
if torch.cuda.is_available():
|
||
allocated = torch.cuda.memory_allocated() / 1024 ** 3
|
||
reserved = torch.cuda.memory_reserved() / 1024 ** 3
|
||
max_allocated = torch.cuda.max_memory_allocated() / 1024 ** 3
|
||
print(f"📊 GPU显存: 已分配 {allocated:.2f}GB | 已预留 {reserved:.2f}GB | 峰值 {max_allocated:.2f}GB")
|
||
|
||
|
||
def format_time(seconds):
|
||
"""格式化时间显示"""
|
||
if seconds < 60:
|
||
return f"{seconds:.0f}秒"
|
||
elif seconds < 3600:
|
||
return f"{seconds / 60:.1f}分钟"
|
||
else:
|
||
return f"{seconds / 3600:.2f}小时"
|
||
|
||
|
||
# ==================== 4. 数据加载 ====================
|
||
def load_data(file_path):
|
||
"""加载数据文件"""
|
||
try:
|
||
if file_path.endswith('.csv'):
|
||
df = pd.read_csv(file_path)
|
||
elif file_path.endswith(('.xlsx', '.xls')):
|
||
df = pd.read_excel(file_path)
|
||
else:
|
||
raise ValueError("不支持的文件格式，请使用 .csv 或 .xlsx")
|
||
|
||
print(f"📄 原始数据列: {df.columns.tolist()}")
|
||
|
||
# 自动识别文本列
|
||
if 'sentence' not in df.columns:
|
||
possible_text_cols = ['sentence', 'content', 'name', 'product_name',
|
||
'goods_name', 'title', 'Title', '商品名称', '商品全称']
|
||
for col in possible_text_cols:
|
||
if col in df.columns:
|
||
df.rename(columns={col: 'sentence'}, inplace=True)
|
||
break
|
||
|
||
# 自动识别标签列
|
||
if 'label' not in df.columns:
|
||
possible_label_cols = ['label', 'class', 'type', 'cate',
|
||
'Type', '分类', '类别']
|
||
for col in possible_label_cols:
|
||
if col in df.columns:
|
||
df.rename(columns={col: 'label'}, inplace=True)
|
||
break
|
||
|
||
assert 'sentence' in df.columns and 'label' in df.columns, \
|
||
"数据必须包含'sentence'(文本)和'label'(标签)列"
|
||
|
||
df = df.dropna(subset=['sentence', 'label']).reset_index(drop=True)
|
||
|
||
print(f"✅ 数据加载成功 | 样本量: {len(df)} | 分类数: {df['label'].nunique()}")
|
||
return df
|
||
except Exception as e:
|
||
warnings.warn(f"❌ 数据加载失败: {str(e)}")
|
||
raise
|
||
|
||
|
||
# ==================== 5. 数据清洗 ====================
|
||
def clean_chinese_text(text):
|
||
"""优化版清洗：保留中文字符、英文字母、数字、空格和重要符号"""
|
||
if not isinstance(text, str):
|
||
return ""
|
||
|
||
cleaned = re.sub(r'[^\u4e00-\u9fa5a-zA-Z0-9\s\-/\.\(\)\（\）\·]', ' ', text)
|
||
cleaned = re.sub(r'\s+', ' ', cleaned).strip()
|
||
return cleaned
|
||
|
||
|
||
def clean_batch_texts(texts, show_sample=True):
|
||
"""批量清洗文本"""
|
||
print("🧼 开始清洗文本数据...")
|
||
cleaned = [clean_chinese_text(t) for t in texts]
|
||
|
||
original_len = sum(len(t) for t in texts)
|
||
cleaned_len = sum(len(t) for t in cleaned)
|
||
if original_len > 0:
|
||
print(f" 原始总字符: {original_len:,} | 清洗后总字符: {cleaned_len:,} | 缩减: {(1 - cleaned_len / original_len) * 100:.1f}%")
|
||
|
||
if show_sample and len(texts) > 0:
|
||
print("\n📝 清洗样例:")
|
||
for i in range(min(3, len(texts))):
|
||
print(f" 原始: {texts[i][:50]}{'...' if len(texts[i]) > 50 else ''}")
|
||
print(f" 清洗: {cleaned[i][:50]}{'...' if len(cleaned[i]) > 50 else ''}")
|
||
print()
|
||
|
||
return cleaned
|
||
|
||
|
||
# ==================== 6. 数据集类 ====================
|
||
class TextDataset(Dataset):
|
||
"""动态编码数据集"""
|
||
|
||
def __init__(self, texts, labels, tokenizer, max_length):
|
||
self.texts = texts
|
||
self.labels = labels
|
||
self.tokenizer = tokenizer
|
||
self.max_length = max_length
|
||
|
||
def __len__(self):
|
||
return len(self.texts)
|
||
|
||
def __getitem__(self, idx):
|
||
encoding = self.tokenizer(
|
||
self.texts[idx],
|
||
max_length=self.max_length,
|
||
padding=False,
|
||
truncation=True,
|
||
return_tensors=None
|
||
)
|
||
return {
|
||
"input_ids": encoding["input_ids"],
|
||
"attention_mask": encoding["attention_mask"],
|
||
"labels": self.labels[idx]
|
||
}
|
||
|
||
|
||
# ==================== 7. Collate函数 ====================
|
||
def collate_fn(batch, tokenizer):
|
||
"""动态填充批处理"""
|
||
input_ids = [item["input_ids"] for item in batch]
|
||
attention_mask = [item["attention_mask"] for item in batch]
|
||
labels = [item["labels"] for item in batch]
|
||
|
||
padded = tokenizer.pad(
|
||
{"input_ids": input_ids, "attention_mask": attention_mask},
|
||
padding=True,
|
||
return_tensors="pt"
|
||
)
|
||
|
||
return {
|
||
"input_ids": padded["input_ids"],
|
||
"attention_mask": padded["attention_mask"],
|
||
"labels": torch.tensor(labels, dtype=torch.long)
|
||
}
|
||
|
||
|
||
# ==================== 8. 评估指标 ====================
|
||
def compute_metrics(eval_pred):
|
||
"""计算多种评估指标"""
|
||
logits, labels = eval_pred
|
||
predictions = np.argmax(logits, axis=-1)
|
||
|
||
return {
|
||
"accuracy": accuracy_score(labels, predictions),
|
||
"f1_macro": f1_score(labels, predictions, average="macro"),
|
||
"f1_weighted": f1_score(labels, predictions, average="weighted"),
|
||
"f1_micro": f1_score(labels, predictions, average="micro"),
|
||
}
|
||
|
||
|
||
# ==================== 9. 类别权重 ====================
|
||
def compute_class_weights(labels, num_classes):
|
||
"""计算类别权重"""
|
||
classes = np.unique(labels)
|
||
weights = compute_class_weight("balanced", classes=classes, y=labels)
|
||
|
||
weight_tensor = torch.ones(num_classes, dtype=torch.float)
|
||
for cls, weight in zip(classes, weights):
|
||
weight_tensor[cls] = weight
|
||
|
||
return weight_tensor
|
||
|
||
|
||
# ==================== 10. 带权重损失的自定义模型 ====================
|
||
class CustomClassificationModel(torch.nn.Module):
|
||
"""包装原始模型，支持自定义损失函数"""
|
||
|
||
def __init__(self, base_model, class_weights=None):
|
||
super().__init__()
|
||
self.base_model = base_model
|
||
self.num_labels = base_model.config.num_labels
|
||
self.config = base_model.config
|
||
|
||
if class_weights is not None:
|
||
self.loss_fct = torch.nn.CrossEntropyLoss(weight=class_weights)
|
||
else:
|
||
self.loss_fct = torch.nn.CrossEntropyLoss()
|
||
|
||
def forward(self, input_ids=None, attention_mask=None, labels=None, **kwargs):
|
||
outputs = self.base_model(
|
||
input_ids=input_ids,
|
||
attention_mask=attention_mask,
|
||
**kwargs
|
||
)
|
||
|
||
logits = outputs.logits
|
||
|
||
loss = None
|
||
if labels is not None:
|
||
loss = self.loss_fct(logits, labels)
|
||
|
||
return SequenceClassifierOutput(
|
||
loss=loss,
|
||
logits=logits,
|
||
hidden_states=outputs.hidden_states if hasattr(outputs, 'hidden_states') else None,
|
||
attentions=outputs.attentions if hasattr(outputs, 'attentions') else None
|
||
)
|
||
|
||
|
||
# ==================== 11. 自定义保存回调（修复非连续张量问题） ====================
|
||
class CustomSaveCallback(TrainerCallback):
|
||
"""
|
||
自定义保存回调：在保存前修复非连续张量
|
||
"""
|
||
def on_save(self, args, state, control, model, tokenizer=None, **kwargs):
|
||
if state.is_world_process_zero:
|
||
# 获取当前保存路径
|
||
checkpoint_dir = os.path.join(args.output_dir, f"checkpoint-{state.global_step}")
|
||
|
||
# 如果是包装模型，获取基础模型
|
||
if hasattr(model, 'base_model'):
|
||
model_to_save = model.base_model
|
||
else:
|
||
model_to_save = model
|
||
|
||
# 确保所有参数连续
|
||
for param in model_to_save.parameters():
|
||
if not param.is_contiguous():
|
||
param.data = param.data.contiguous()
|
||
|
||
# 保存模型（使用pickle格式）
|
||
try:
|
||
model_to_save.save_pretrained(checkpoint_dir)
|
||
if tokenizer:
|
||
tokenizer.save_pretrained(checkpoint_dir)
|
||
print(f"✅ 检查点已保存: {checkpoint_dir}")
|
||
except Exception as e:
|
||
print(f"⚠️ 保存检查点失败: {e}")
|
||
# 备用方案：手动保存
|
||
self._manual_save(model_to_save, tokenizer, checkpoint_dir)
|
||
|
||
def _manual_save(self, model, tokenizer, save_dir):
|
||
"""手动保存模型"""
|
||
try:
|
||
os.makedirs(save_dir, exist_ok=True)
|
||
model.config.save_pretrained(save_dir)
|
||
|
||
state_dict = model.state_dict()
|
||
contiguous_state_dict = {}
|
||
for key, tensor in state_dict.items():
|
||
if not tensor.is_contiguous():
|
||
contiguous_state_dict[key] = tensor.contiguous()
|
||
else:
|
||
contiguous_state_dict[key] = tensor
|
||
|
||
torch.save(contiguous_state_dict, os.path.join(save_dir, "pytorch_model.bin"))
|
||
if tokenizer:
|
||
tokenizer.save_pretrained(save_dir)
|
||
print(f"✅ 检查点已保存 (手动): {save_dir}")
|
||
except Exception as e:
|
||
print(f"❌ 保存失败: {e}")
|
||
|
||
|
||
# ==================== 12. TensorBoard回调 ====================
|
||
class TensorBoardCallback(TrainerCallback):
|
||
"""TensorBoard回调"""
|
||
|
||
def __init__(self, log_dir):
|
||
self.writer = SummaryWriter(log_dir)
|
||
|
||
def on_log(self, args, state, control, logs=None, **kwargs):
|
||
if logs:
|
||
step = state.global_step
|
||
for key, value in logs.items():
|
||
if isinstance(value, (int, float)):
|
||
self.writer.add_scalar(key, value, step)
|
||
|
||
def on_train_end(self, args, state, control, **kwargs):
|
||
self.writer.close()
|
||
|
||
|
||
# ==================== 13. 模型初始化 ====================
|
||
def init_model(num_labels, class_weights=None):
|
||
"""初始化模型和分词器"""
|
||
model_path = get_model_path(
|
||
Config.MODEL_ID,
|
||
Config.CACHE_DIR,
|
||
Config.USE_MODELSCOPE
|
||
)
|
||
|
||
print(f"\n🔄 加载模型: {model_path}")
|
||
start_time = time.time()
|
||
|
||
tokenizer = AutoTokenizer.from_pretrained(
|
||
model_path,
|
||
trust_remote_code=True
|
||
)
|
||
|
||
base_model = AutoModelForSequenceClassification.from_pretrained(
|
||
model_path,
|
||
num_labels=num_labels,
|
||
trust_remote_code=True
|
||
)
|
||
|
||
# 确保模型参数是连续的（修复非连续张量问题）
|
||
base_model = base_model.to(Config.DEVICE)
|
||
|
||
# 遍历所有参数，确保连续
|
||
for param in base_model.parameters():
|
||
if not param.is_contiguous():
|
||
param.data = param.data.contiguous()
|
||
|
||
if class_weights is not None:
|
||
class_weights = class_weights.to(Config.DEVICE)
|
||
model = CustomClassificationModel(base_model, class_weights)
|
||
else:
|
||
model = base_model
|
||
|
||
load_time = time.time() - start_time
|
||
print(f"✅ 模型加载完成，耗时: {format_time(load_time)}")
|
||
|
||
total_params = sum(p.numel() for p in model.parameters())
|
||
trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
|
||
print(f"📊 模型参数总量: {total_params:,} ({total_params / 1e6:.1f}M)")
|
||
print(f"📊 可训练参数: {trainable_params:,} ({trainable_params / 1e6:.1f}M)")
|
||
|
||
return tokenizer, model
|
||
|
||
|
||
# ==================== 14. 训练配置 ====================
|
||
def get_training_args():
|
||
"""获取训练参数配置"""
|
||
return TrainingArguments(
|
||
output_dir=Config.OUTPUT_DIR,
|
||
num_train_epochs=Config.NUM_EPOCHS,
|
||
per_device_train_batch_size=Config.BATCH_SIZE,
|
||
per_device_eval_batch_size=Config.BATCH_SIZE * 2,
|
||
gradient_accumulation_steps=Config.GRADIENT_ACCUMULATION_STEPS,
|
||
learning_rate=Config.LEARNING_RATE,
|
||
warmup_ratio=Config.WARMUP_RATIO,
|
||
weight_decay=Config.WEIGHT_DECAY,
|
||
|
||
# 日志和保存
|
||
logging_dir=Config.LOG_DIR,
|
||
logging_steps=Config.LOGGING_STEPS,
|
||
logging_first_step=True,
|
||
eval_strategy="steps",
|
||
eval_steps=Config.EVAL_SAVE_STEPS,
|
||
save_strategy="steps",
|
||
save_steps=Config.EVAL_SAVE_STEPS,
|
||
save_total_limit=Config.SAVE_TOTAL_LIMIT,
|
||
load_best_model_at_end=True,
|
||
metric_for_best_model=Config.METRIC_FOR_BEST_MODEL,
|
||
greater_is_better=Config.GREATER_IS_BETTER,
|
||
|
||
# 性能优化
|
||
fp16=Config.FP16,
|
||
dataloader_num_workers=Config.NUM_WORKERS,
|
||
dataloader_pin_memory=True,
|
||
|
||
# 其他
|
||
report_to="tensorboard",
|
||
seed=Config.SEED,
|
||
push_to_hub=False,
|
||
lr_scheduler_type="linear",
|
||
remove_unused_columns=False,
|
||
dataloader_drop_last=False,
|
||
|
||
# ✅ 关键修复：禁用 safetensors
|
||
save_safetensors=False,
|
||
)
|
||
|
||
|
||
# ==================== 15. 保存模型 ====================
|
||
def save_model(model, tokenizer, save_dir):
|
||
"""保存模型"""
|
||
print(f"\n💾 保存模型到: {save_dir}")
|
||
os.makedirs(save_dir, exist_ok=True)
|
||
|
||
if hasattr(model, 'base_model'):
|
||
model_to_save = model.base_model
|
||
else:
|
||
model_to_save = model
|
||
|
||
# 确保所有参数连续
|
||
for param in model_to_save.parameters():
|
||
if not param.is_contiguous():
|
||
param.data = param.data.contiguous()
|
||
|
||
try:
|
||
model_to_save.save_pretrained(save_dir)
|
||
tokenizer.save_pretrained(save_dir)
|
||
print(f"✅ 模型已保存到: {save_dir}")
|
||
return True
|
||
except Exception as e:
|
||
print(f"⚠️ 保存失败: {e}")
|
||
print(" 尝试使用手动方式保存...")
|
||
|
||
# 备用方案：手动保存
|
||
try:
|
||
model_to_save.config.save_pretrained(save_dir)
|
||
|
||
state_dict = model_to_save.state_dict()
|
||
contiguous_state_dict = {}
|
||
for key, tensor in state_dict.items():
|
||
if not tensor.is_contiguous():
|
||
contiguous_state_dict[key] = tensor.contiguous()
|
||
else:
|
||
contiguous_state_dict[key] = tensor
|
||
|
||
torch.save(contiguous_state_dict, os.path.join(save_dir, "pytorch_model.bin"))
|
||
tokenizer.save_pretrained(save_dir)
|
||
print(f"✅ 模型已保存到: {save_dir}")
|
||
return True
|
||
except Exception as e:
|
||
print(f"❌ 保存失败: {e}")
|
||
return False
|
||
|
||
|
||
# ==================== 16. 预测函数 ====================
|
||
@torch.no_grad()
|
||
def batch_predict(texts, model, tokenizer, label_map, top_k=3, batch_size=64):
|
||
"""批量预测"""
|
||
model.eval()
|
||
all_results = []
|
||
|
||
cleaned_texts = [clean_chinese_text(t) for t in texts]
|
||
|
||
for i in tqdm(range(0, len(cleaned_texts), batch_size), desc="预测中"):
|
||
batch = cleaned_texts[i:i + batch_size]
|
||
inputs = tokenizer(
|
||
batch,
|
||
return_tensors="pt",
|
||
truncation=True,
|
||
padding=True,
|
||
max_length=Config.MAX_LENGTH
|
||
).to(Config.DEVICE)
|
||
|
||
outputs = model(**inputs)
|
||
|
||
if hasattr(outputs, 'logits'):
|
||
logits = outputs.logits
|
||
else:
|
||
logits = outputs
|
||
|
||
probs = torch.softmax(logits, dim=1).cpu()
|
||
|
||
for prob in probs:
|
||
top_probs, top_indices = torch.topk(prob, k=min(top_k, len(label_map)))
|
||
results = [
|
||
{
|
||
"category": label_map[idx.item()],
|
||
"confidence": prob.item()
|
||
}
|
||
for prob, idx in zip(top_probs, top_indices)
|
||
]
|
||
all_results.append(results)
|
||
|
||
return all_results
|
||
|
||
|
||
# ==================== 17. 模型验证 ====================
|
||
def validate_model(model, tokenizer, sample_texts):
|
||
"""快速验证模型"""
|
||
print("\n🔍 模型验证...")
|
||
try:
|
||
inputs = tokenizer(
|
||
sample_texts,
|
||
return_tensors="pt",
|
||
padding=True,
|
||
truncation=True,
|
||
max_length=Config.MAX_LENGTH
|
||
).to(Config.DEVICE)
|
||
|
||
outputs = model(**inputs)
|
||
|
||
if hasattr(outputs, 'logits'):
|
||
logits = outputs.logits
|
||
else:
|
||
logits = outputs
|
||
|
||
print(f"✅ 模型验证通过，输出形状: {logits.shape}")
|
||
return True
|
||
except Exception as e:
|
||
print(f"❌ 模型验证失败: {e}")
|
||
return False
|
||
|
||
|
||
# ==================== 18. 数据统计 ====================
|
||
def print_data_stats(df):
|
||
"""打印数据统计信息"""
|
||
print("\n" + "=" * 50)
|
||
print("📊 数据统计")
|
||
print("=" * 50)
|
||
print(f"总样本量: {len(df)}")
|
||
print(f"分类数: {df['label'].nunique()}")
|
||
|
||
label_counts = df['label'].value_counts()
|
||
print(f"\n📈 类别分布 (Top 10):")
|
||
for label, count in label_counts.head(10).items():
|
||
pct = count / len(df) * 100
|
||
bar = "█" * int(pct / 2)
|
||
print(f" {label}: {count} ({pct:.1f}%) {bar}")
|
||
|
||
tail_threshold = 10
|
||
tail_count = (label_counts <= tail_threshold).sum()
|
||
tail_samples = label_counts[label_counts <= tail_threshold].sum()
|
||
print(f"\n📉 长尾类别 (≤{tail_threshold}条): {tail_count}个类别, 共{tail_samples}条样本")
|
||
|
||
text_lengths = df['sentence'].str.len()
|
||
print(f"\n📏 文本长度统计:")
|
||
print(f" 平均: {text_lengths.mean():.1f}")
|
||
print(f" 中位数: {text_lengths.median():.0f}")
|
||
print(f" 最大: {text_lengths.max()}")
|
||
print(f" 最小: {text_lengths.min()}")
|
||
print("=" * 50 + "\n")
|
||
|
||
|
||
# ==================== 19. 主函数 ====================
|
||
def main():
|
||
"""主训练流程"""
|
||
print("\n" + "=" * 60)
|
||
print("🚀 RoBERTa-wwm-ext-large 商品分类微调 (ModelScope版)")
|
||
print(f"📅 启动时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
|
||
print(f"💻 设备: {Config.DEVICE}")
|
||
print(f"📦 模型来源: {'ModelScope' if Config.USE_MODELSCOPE else 'HuggingFace'}")
|
||
if torch.cuda.is_available():
|
||
print(f"🎮 GPU: {torch.cuda.get_device_name(0)}")
|
||
print(f"💾 显存: {torch.cuda.get_device_properties(0).total_memory / 1024 ** 3:.1f}GB")
|
||
print("=" * 60)
|
||
|
||
total_start = time.time()
|
||
|
||
set_seed(Config.SEED)
|
||
|
||
os.makedirs(Config.OUTPUT_DIR, exist_ok=True)
|
||
os.makedirs(Config.LOG_DIR, exist_ok=True)
|
||
os.makedirs(Config.SAVE_DIR, exist_ok=True)
|
||
os.makedirs(Config.CACHE_DIR, exist_ok=True)
|
||
|
||
# ===== 1. 加载数据 =====
|
||
print(f"\n📂 加载数据: {Config.DATA_FILE}")
|
||
df = load_data(Config.DATA_FILE)
|
||
|
||
# ===== 2. 数据清洗 =====
|
||
df['sentence_cleaned'] = clean_batch_texts(df['sentence'].tolist())
|
||
df = df[df['sentence_cleaned'].str.len() > 0].reset_index(drop=True)
|
||
print(f"✅ 清洗后剩余: {len(df)} 条")
|
||
|
||
# ===== 3. 数据统计 =====
|
||
print_data_stats(df)
|
||
|
||
# ===== 4. 标签编码 =====
|
||
print("🏷️ 处理标签...")
|
||
label_encoder = LabelEncoder()
|
||
df['label_id'] = label_encoder.fit_transform(df['label'])
|
||
label_map = {i: label for i, label in enumerate(label_encoder.classes_)}
|
||
num_labels = len(label_map)
|
||
|
||
print(f"✅ 类别数: {num_labels}")
|
||
print(f" 标签示例: {list(label_map.items())[:5]}")
|
||
|
||
joblib.dump(label_encoder, "label_encoder_roberta_large.pkl")
|
||
print("✅ 标签编码器已保存: label_encoder_roberta_large.pkl")
|
||
|
||
# ===== 5. 类别权重 =====
|
||
print("\n⚖️ 计算类别权重...")
|
||
class_weights = compute_class_weights(df['label_id'].values, num_labels)
|
||
print(f" 权重范围: {class_weights.min():.3f} - {class_weights.max():.3f}")
|
||
print(f" 平均权重: {class_weights.mean():.3f}")
|
||
|
||
# ===== 6. 划分数据集 =====
|
||
print("\n📊 划分数据集...")
|
||
train_df, eval_df = train_test_split(
|
||
df,
|
||
test_size=Config.TEST_SIZE,
|
||
random_state=Config.SEED,
|
||
stratify=df["label_id"]
|
||
)
|
||
print(f" 训练集: {len(train_df)} 条")
|
||
print(f" 验证集: {len(eval_df)} 条")
|
||
|
||
# ===== 7. 初始化模型 =====
|
||
tokenizer, model = init_model(num_labels, class_weights)
|
||
model.to(Config.DEVICE)
|
||
print(f"✅ 模型已加载到: {Config.DEVICE}")
|
||
print_gpu_memory()
|
||
|
||
# ===== 8. 模型验证 =====
|
||
sample_texts = ["九阳电饭煲家用", "苹果手机iPhone", "耐克运动鞋"]
|
||
if not validate_model(model, tokenizer, sample_texts):
|
||
print("⚠️ 模型验证失败，继续训练...")
|
||
|
||
# ===== 9. 准备数据集 =====
|
||
print("\n📦 准备数据集...")
|
||
train_dataset = TextDataset(
|
||
train_df['sentence_cleaned'].tolist(),
|
||
train_df['label_id'].tolist(),
|
||
tokenizer,
|
||
Config.MAX_LENGTH
|
||
)
|
||
eval_dataset = TextDataset(
|
||
eval_df['sentence_cleaned'].tolist(),
|
||
eval_df['label_id'].tolist(),
|
||
tokenizer,
|
||
Config.MAX_LENGTH
|
||
)
|
||
|
||
# ===== 10. 配置训练器 =====
|
||
print("\n⚙️ 配置训练器...")
|
||
training_args = get_training_args()
|
||
|
||
def data_collator(batch):
|
||
return collate_fn(batch, tokenizer)
|
||
|
||
trainer = Trainer(
|
||
model=model,
|
||
args=training_args,
|
||
train_dataset=train_dataset,
|
||
eval_dataset=eval_dataset,
|
||
data_collator=data_collator,
|
||
compute_metrics=compute_metrics,
|
||
callbacks=[
|
||
EarlyStoppingCallback(early_stopping_patience=Config.EARLY_STOPPING_PATIENCE),
|
||
TensorBoardCallback(Config.LOG_DIR),
|
||
CustomSaveCallback(), # ✅ 添加自定义保存回调
|
||
]
|
||
)
|
||
|
||
# ===== 11. 断点续训 =====
|
||
checkpoint_dir = None
|
||
checkpoints = glob.glob(os.path.join(Config.OUTPUT_DIR, "checkpoint-*"))
|
||
if checkpoints:
|
||
latest_checkpoint = max(checkpoints, key=os.path.getctime)
|
||
print(f"📂 发现检查点: {latest_checkpoint}")
|
||
checkpoint_dir = latest_checkpoint
|
||
|
||
# ===== 12. 开始训练 =====
|
||
print("\n" + "=" * 60)
|
||
print("🚀 开始训练")
|
||
print("=" * 60)
|
||
|
||
train_start = time.time()
|
||
|
||
if checkpoint_dir:
|
||
trainer.train(resume_from_checkpoint=checkpoint_dir)
|
||
else:
|
||
trainer.train()
|
||
|
||
train_time = time.time() - train_start
|
||
print(f"\n✅ 训练完成，耗时: {format_time(train_time)}")
|
||
|
||
# ===== 13. 保存模型 =====
|
||
save_model(model, tokenizer, Config.SAVE_DIR)
|
||
|
||
# ===== 14. 最终评估 =====
|
||
print("\n📊 最终评估...")
|
||
eval_results = trainer.evaluate()
|
||
print(f" Eval Loss: {eval_results['eval_loss']:.4f}")
|
||
print(f" Eval Accuracy: {eval_results.get('eval_accuracy', 0):.4f}")
|
||
print(f" Eval F1 Macro: {eval_results.get('eval_f1_macro', 0):.4f}")
|
||
print(f" Eval F1 Weighted: {eval_results.get('eval_f1_weighted', 0):.4f}")
|
||
|
||
# ===== 15. 推理测试 =====
|
||
print("\n🧪 推理测试...")
|
||
test_samples = [
|
||
"九阳3升家用多功能电饭煲F30S-S160",
|
||
"苹果iPhone 15 Pro Max 512GB 黑色钛金属",
|
||
"海尔双开门冰箱BCD-600WGHSS19B8U1",
|
||
"Lining李宁女款长裤 AKYV092-1 黑色",
|
||
"珀莱雅双抗精华美白特证版补水保湿套装50ml礼盒",
|
||
]
|
||
|
||
predictions = batch_predict(test_samples, model, tokenizer, label_map, top_k=3)
|
||
|
||
print("\n📋 预测结果:")
|
||
print("-" * 80)
|
||
for sample, preds in zip(test_samples, predictions):
|
||
print(f"输入: {sample}")
|
||
for i, p in enumerate(preds):
|
||
print(f" Top{i + 1}: {p['category']} (置信度: {p['confidence']:.4f})")
|
||
print("-" * 80)
|
||
|
||
# ===== 16. 总结 =====
|
||
total_time = time.time() - total_start
|
||
print("\n" + "=" * 60)
|
||
print("🎉 训练完成!")
|
||
print(f"📅 结束时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
|
||
print(f"⏱️ 总耗时: {format_time(total_time)}")
|
||
print(f"⏱️ 训练耗时: {format_time(train_time)}")
|
||
print(f"📊 最佳模型指标: {Config.METRIC_FOR_BEST_MODEL} = {trainer.state.best_metric:.4f}")
|
||
print("=" * 60)
|
||
|
||
|
||
# ==================== 20. 入口 ====================
|
||
if __name__ == "__main__":
|
||
main() |