python
复制代码
"""
实验3: TextRNN (BiLSTM + Attention) 文本分类 (exp3)
====================================================
"""
import os, json, time, random, warnings
warnings.filterwarnings('ignore')
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
from sklearn.metrics import f1_score, accuracy_score, classification_report
BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
DATA_DIR = os.path.join(BASE_DIR, 'data')
LOGS_DIR = os.path.join(BASE_DIR, 'logs')
MODELS_DIR = os.path.join(BASE_DIR, 'models')
SUBMISSIONS_DIR = os.path.join(BASE_DIR, 'submissions')
for d in [LOGS_DIR, MODELS_DIR, SUBMISSIONS_DIR]:
os.makedirs(d, exist_ok=True)
LABEL_MAP = {0: '科技', 1: '股票', 2: '体育', 3: '娱乐', 4: '时政',
5: '社会', 6: '教育', 7: '财经', 8: '家居', 9: '游戏',
10: '房产', 11: '时尚', 12: '彩票', 13: '星座'}
SEED = 42
MAX_LEN = 256
VOCAB_SIZE = 7550
EMBED_DIM = 100
# 【核心参数大换血】
HIDDEN_SIZE = 128 # LSTM 单向隐藏层维度。因为是双向(Bi)，最终特征会变成 128 * 2 = 256 维。
NUM_LAYERS = 2 # 堆叠两层 LSTM，让模型能学到更深层的逻辑。
DROPOUT = 0.15
NUM_CLASSES = 14
BATCH_SIZE = 16 # 【注意】RNN 极其吃显存，所以 Batch Size 从 128 降到了 16！
LR = 2e-4
EPOCHS = 10
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
def set_seed(seed):
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
set_seed(SEED)
# ==========================================
# 关卡一：带有"真实长度"记录仪的数据加工厂
# ==========================================
class NewsDataset(Dataset):
def __init__(self, texts, labels=None, max_len=MAX_LEN):
# 【招募库管员】
# 把所有新闻文本（长串的数字字符串）和标签存进仓库
self.texts = texts
self.labels = labels
self.max_len = max_len # 规定流水线上的标准包装盒大小（256）
def __len__(self):
# 【清点库存】
# 告诉 DataLoader，我们仓库里总共有多少篇新闻
return len(self.texts)
def __getitem__(self, idx):
# 【核心加工流水线】：当 DataLoader 来要第 idx 篇新闻时，现场加工
# 1. 拆解零件：把 "345 889 12" 这种整块的字符串，切成一个个整数组成的列表
tokens = [int(t) for t in self.texts[idx].split()]
# ---------------------------------------------------------
# 🌟【全场最核心的改动：记录真实长度】🌟
# ---------------------------------------------------------
# 为什么用 min？
# 如果句子长度是 50，min(50, 256) = 50。真实长度就是 50。
# 如果句子长度是 300，超过了盒子大小，一会儿要被切掉，所以它的有效真实长度最多只能算作 256。
# 【面试防坑】：这个 length 极其重要！它要被送给后方的 LSTM，告诉它"算到这个长度就赶紧停下，后面的全是 0，不用看了！"
length = min(len(tokens), self.max_len)
# todo：GPU 是一个极其死板的"流水线工人"，它根本不懂什么是"句子"，它只认识方方正正的矩阵
# 所以固定长度
# 2. 截断与填充 (把残次品变成标准件)
if len(tokens) > self.max_len:
# 太长的，一刀切掉尾巴
tokens = tokens[:self.max_len]
else:
# 太短的，屁股后面全塞上 0 (Padding)
tokens = tokens + [0] * (self.max_len - len(tokens))
# 3. 装车打包：转换成 PyTorch 喜欢的 Tensor (长整型)矩阵
x = torch.tensor(tokens, dtype=torch.long)
# 4. 发货：带着"小尾巴"出厂
if self.labels is not None:
# 训练集：不仅返回标准件 x，返回标签 y，还把【真实长度 length】一起打包发走！
return x, torch.tensor(self.labels[idx], dtype=torch.long), length
# 测试集：没有标签，但依然要带上【真实长度 length】
return x, length
# ==========================================
# 关卡三：划重点的神器 (Attention 机制)
# ==========================================
class Attention(nn.Module):
def __init__(self, hidden_size):
super().__init__()
# 【招募阅卷老师】
# hidden_size 在这里是 256 (因为是双向 LSTM，128 * 2)
# 1. 投影矩阵 W：相当于给阅卷老师配了一副"透视眼镜"
self.W = nn.Linear(hidden_size, hidden_size)
# 2. 核心评价向量 v：相当于阅卷老师脑子里的"标准答案"
# nn.Parameter 意味着这是一个可以随着训练不断优化的参数，模型会自己学到什么特征最重要。
self.v = nn.Parameter(torch.randn(hidden_size))
def forward(self, h, mask=None):
# 【开始阅卷】
# 假设当前输入 h (Bi-LSTM的输出) 的形状是: [16, 256, 256]
# 代表 16句话，每句话 256个词，每个词已经被提炼成了 256维 的高级特征。
# ---------------------------------------------------------
# 第一步：给每个词打初评分
# ---------------------------------------------------------
# 先让 256 维的特征戴上透视眼镜 (W)，然后用 tanh 激活函数把数值压缩到 -1 到 1 之间
score = torch.tanh(self.W(h))
# 让处理过的特征，去和老师脑子里的标准答案 (v) 算内积 (matmul)
# 内积越大，说明这个词的特征和标准答案越匹配，分数越高！
# 形状突变: [16, 256, 256] 和 [256] 相乘 -> [16, 256] (每句话里的 256 个词，各自得到了一个具体的分数)
score = torch.matmul(score, self.v)
# ---------------------------------------------------------
# 第二步：极其冷酷的 Mask (遮蔽废话)
# ---------------------------------------------------------
if mask is not None:
# 面试官必考：为什么要填 -1e9 (负十亿)？
# 答：因为后面要过 Softmax。补齐的 0 没有任何语义，如果初评分数是 0，过 Softmax 后依然会分走一部分权重。
# 把补齐位置的分数强行改成 -1e9，Softmax(-1e9) 会无限趋近于绝对的 0。彻底封杀废话的发言权！
score = score.masked_fill(~mask, -1e9)
# ---------------------------------------------------------
# 第三步：转化为重要性百分比 (权重分配)
# ---------------------------------------------------------
# softmax 会把 256 个词的分数，变成加起来等于 100% 的概率分布。
# 比如"科技"占 80%，"发布"占 15%，标点符号占 0%。
# unsqueeze(2) 是为了补全维度，方便后面做乘法。形状突变: [16, 256] -> [16, 256, 1]
attn = torch.softmax(score, dim=1).unsqueeze(2)
# ---------------------------------------------------------
# 第四步：加权融合 (提取终极精华)
# ---------------------------------------------------------
# h [16, 256, 256] * attn [16, 256, 1]
# 这一步极其优美：把每个词的 256 维特征，乘以它分配到的权重。废话乘以 0 直接消失，关键词乘以大比重被无限放大。
# 最后用 sum(dim=1) 把这 256 个加权后的词压缩融合成 1 个向量！
# 形状突变: [16, 256, 256] -> [16, 256]
out = (h * attn).sum(dim=1)
return out # 吐出这 16 句话的终极浓缩精华！
# ==========================================
# 终极缝合怪：TextRNN (BiLSTM + Attention) 主干网络
# ==========================================
class TextRNN(nn.Module):
def __init__(self):
super().__init__()
# 1. 词嵌入层 (查字典)：7550个词，每个词变100维向量。遇到0(补齐位)直接输出全0。
self.embedding = nn.Embedding(VOCAB_SIZE, EMBED_DIM, padding_idx=0)
# 2. 核心大闸：双向 LSTM (Bi-LSTM)
# EMBED_DIM: 输入维度(100)
# HIDDEN_SIZE: 隐藏层维度(128)。因为开了双向(bidirectional=True)，输出维度会翻倍变成 256。
# batch_first=True: 告诉 PyTorch 输入的第一维是 batch_size [batch, seq_len, feature]
self.lstm = nn.LSTM(EMBED_DIM, HIDDEN_SIZE, NUM_LAYERS,
batch_first=True, bidirectional=True,
dropout=DROPOUT if NUM_LAYERS > 1 else 0)
# 3. 注意力机制 (划重点)：输入是 Bi-LSTM 的输出维度 (128 * 2 = 256)
self.attention = Attention(HIDDEN_SIZE * 2)
# 4. 随机失活防过拟合
self.dropout = nn.Dropout(DROPOUT)
# 5. 最终法官 (全连接层)：把提炼出的 256维 精华，映射到 14个 新闻类别上打分
self.fc = nn.Linear(HIDDEN_SIZE * 2, NUM_CLASSES)
def forward(self, x, lengths=None):
# 假设 x 的形状: [16, 256] (16句话，每句强行Pad到了256长)
# 【极其关键的 Mask 掩码生成】
# 找出 x 里不是 0 的真实单词位置，标为 True；是 0 的补齐位，标为 False。
# 这个 mask 后面会送给 Attention，用来屏蔽掉补齐位的注意力打分。
mask = (x != 0)
# 查字典：形状突变 [16, 256] -> [16, 256, 100]
x = self.embedding(x)
x = self.dropout(x)
# 【面试高光时刻：动态序列打包】
if lengths is not None:
# 第一步：真空压缩 (剥离所有的 Padding 0)
# 把原本方方正正带有海量 0 的矩阵，按真实长度压缩成极其紧凑的 PackedSequence 对象。
packed = nn.utils.rnn.pack_padded_sequence(
x, lengths.cpu().clamp(min=1), batch_first=True, enforce_sorted=False)
# 第二步：纯净版 LSTM 计算 (速度极快，且记忆不会被 0 污染)
out, _ = self.lstm(packed)
# 第三步：充气解压 (把压缩包还原成方阵)
# 此时的 out 恢复成了 [16, 256, 256]。原来补 0 的位置，现在的特征依然是 0。
out, _ = nn.utils.rnn.pad_packed_sequence(out, batch_first=True, total_length=mask.size(1))
else:
# 裸考模式：如果没有提供真实长度，只能带 0 傻算（慢且效果差）
out, _ = self.lstm(x)
# 经过 LSTM 后，out 形状: [16, 256, 256]
# 把带有 256 步特征的 out 和刚才生成的 mask 交给 Attention 老师划重点。
# Attention 老师会把这 256 个词的特征，按重要性加权浓缩成 1 个特征！
# 形状突变: [16, 256, 256] -> [16, 256]
out = self.attention(out, mask)
out = self.dropout(out)
# 最终审判：输出 14 个类别的原始打分 (Logits)
# 形状突变: [16, 256] -> [16, 14]
return self.fc(out)
# ==========================================
# 训练引擎与评估体系 (与 TextCNN 逻辑一致，但加入了 lengths)
# ==========================================
def get_class_weights(labels):
# 【劫富济贫】：统计各类别频次并取倒数，给样本极少的类（如星座）极大的 Loss 权重惩罚。
counts = np.bincount(labels, minlength=NUM_CLASSES)
weights = 1.0 / (counts + 1e-6)
weights = weights / weights.sum() * NUM_CLASSES # 均值归一化，防止 Loss 整体过小
return torch.tensor(weights, dtype=torch.float32)
def train_epoch(model, loader, optimizer, criterion, device):
model.train() # 开启训练模式 (启用 Dropout)
total_loss = 0
all_preds, all_labels = [], []
for batch in loader:
# 【区别点】：除了 x 和 y，还要把句子的真实长度 lengths 取出来
x, y, lengths = batch[0].to(device), batch[1].to(device), batch[2]
optimizer.zero_grad() # 1. 洗杯子 (梯度清零)
# 2. 前向传播：带着 lengths 一起送进模型，触发 pack_padded_sequence 加速
logits = model(x, lengths)
loss = criterion(logits, y) # 3. 算误差 (已包含类别加权)
loss.backward() # 4. 反向求导 (计算梯度)
# 【面试防坑】：梯度裁剪，防止梯度爆炸毁掉网络参数
nn.utils.clip_grad_norm_(model.parameters(), 5.0)
optimizer.step() # 5. 更新权重参数
total_loss += loss.item() * x.size(0)
# 取出 14 个打分里最高的那一类的索引 (0~13) 作为预测结果
all_preds.extend(logits.argmax(1).cpu().numpy())
all_labels.extend(y.cpu().numpy())
# 返回平均 Loss 和抗数据不平衡的 Macro-F1 分数
return total_loss / len(loader.dataset), f1_score(all_labels, all_preds, average='macro')
@torch.no_grad() # 关闭底层梯度追踪，极大节省显存
def eval_epoch(model, loader, criterion, device):
model.eval() # 开启验证模式 (关闭 Dropout，保证预测稳定)
total_loss = 0
all_preds, all_labels = [], []
for batch in loader:
x, y, lengths = batch[0].to(device), batch[1].to(device), batch[2]
logits = model(x, lengths)
loss = criterion(logits, y)
total_loss += loss.item() * x.size(0)
all_preds.extend(logits.argmax(1).cpu().numpy())
all_labels.extend(y.cpu().numpy())
f1 = f1_score(all_labels, all_preds, average='macro')
acc = accuracy_score(all_labels, all_preds)
return total_loss / len(loader.dataset), f1, acc, all_preds, all_labels
@torch.no_grad()
def predict(model, loader, device):
model.eval()
all_preds = []
for batch in loader:
x = batch[0].to(device)
# 兼容性安检：测试集可能不带标签，长度数据可能在索引 1 也可能在 2
lengths = batch[1] if len(batch) == 2 else batch[2]
logits = model(x, lengths)
all_preds.extend(logits.argmax(1).cpu().numpy())
return all_preds
def main():
print("=" * 60)
print(f"实验3: TextRNN BiLSTM+Attention (max_len={MAX_LEN}, device={DEVICE})")
print("=" * 60)
train_df = pd.read_csv(os.path.join(DATA_DIR, 'train_split.csv'), sep='\t')
val_df = pd.read_csv(os.path.join(DATA_DIR, 'val_split.csv'), sep='\t')
test_df = pd.read_csv(os.path.join(DATA_DIR, 'test_a.csv'), sep='\t')
train_ds = NewsDataset(train_df['text'].tolist(), train_df['label'].tolist())
val_ds = NewsDataset(val_df['text'].tolist(), val_df['label'].tolist())
test_ds = NewsDataset(test_df['text'].tolist())
train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True, num_workers=4, pin_memory=True)
val_loader = DataLoader(val_ds, batch_size=BATCH_SIZE*2, shuffle=False, num_workers=4, pin_memory=True)
test_loader = DataLoader(test_ds, batch_size=BATCH_SIZE*2, shuffle=False, num_workers=4, pin_memory=True)
model = TextRNN().to(DEVICE)
param_count = sum(p.numel() for p in model.parameters())
print(f"模型参数量: {param_count/1e6:.2f}M")
class_weights = get_class_weights(train_df['label'].values).to(DEVICE)
criterion = nn.CrossEntropyLoss(weight=class_weights)
optimizer = torch.optim.Adam(model.parameters(), lr=LR)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=EPOCHS)
best_f1 = 0
history = []
for epoch in range(1, EPOCHS + 1):
t0 = time.time()
train_loss, train_f1 = train_epoch(model, train_loader, optimizer, criterion, DEVICE)
val_loss, val_f1, val_acc, val_preds, val_labels = eval_epoch(model, val_loader, criterion, DEVICE)
scheduler.step()
elapsed = time.time() - t0
print(f"Epoch {epoch:2d}/{EPOCHS} | "
f"train_loss={train_loss:.4f} train_f1={train_f1:.4f} | "
f"val_loss={val_loss:.4f} val_f1={val_f1:.4f} val_acc={val_acc:.4f} | "
f"time={elapsed:.1f}s")
history.append({
'epoch': epoch, 'train_loss': train_loss, 'train_f1': train_f1,
'val_loss': val_loss, 'val_f1': val_f1, 'val_acc': val_acc
})
if val_f1 > best_f1:
best_f1 = val_f1
torch.save(model.state_dict(), os.path.join(MODELS_DIR, 'textrnn_best.pt'))
print(f" -> 保存最佳模型 (val_f1={best_f1:.4f})")
model.load_state_dict(torch.load(os.path.join(MODELS_DIR, 'textrnn_best.pt')))
_, final_f1, final_acc, val_preds, val_labels = eval_epoch(model, val_loader, criterion, DEVICE)
per_class = f1_score(val_labels, val_preds, average=None)
per_class_dict = {LABEL_MAP[i]: float(f'{v:.4f}') for i, v in enumerate(per_class)}
print(f"\n最终验证集: macro-F1={final_f1:.4f}, accuracy={final_acc:.4f}")
print(classification_report(val_labels, val_preds,
target_names=[LABEL_MAP[i] for i in range(14)], digits=4))
# 保存val logits
model.eval()
val_logits_list = []
with torch.no_grad():
for batch in val_loader:
x, _, lengths = batch[0].to(DEVICE), batch[1], batch[2]
logits = model(x, lengths)
val_logits_list.append(logits.cpu())
val_logits_np = torch.cat(val_logits_list, dim=0).numpy()
np.save(os.path.join(MODELS_DIR, 'textrnn_val_logits.npy'), val_logits_np)
# 保存test logits + 预测
test_logits_list = []
with torch.no_grad():
for batch in test_loader:
x = batch[0].to(DEVICE)
lengths = batch[1] if len(batch) == 2 else batch[2]
logits = model(x, lengths)
test_logits_list.append(logits.cpu())
test_logits_np = torch.cat(test_logits_list, dim=0).numpy()
np.save(os.path.join(MODELS_DIR, 'textrnn_test_logits.npy'), test_logits_np)
test_preds = test_logits_np.argmax(axis=1)
sub_df = pd.DataFrame({'label': test_preds})
sub_df.to_csv(os.path.join(SUBMISSIONS_DIR, 'submission_textrnn.csv'), index=False)
print(f"-> logits已保存: textrnn_val_logits.npy, textrnn_test_logits.npy")
result = {
'model': 'TextRNN_BiLSTM_Attention',
'max_len': MAX_LEN, 'embed_dim': EMBED_DIM,
'hidden_size': HIDDEN_SIZE, 'num_layers': NUM_LAYERS,
'batch_size': BATCH_SIZE, 'lr': LR, 'epochs': EPOCHS,
'best_val_f1': float(f'{best_f1:.4f}'),
'best_val_acc': float(f'{final_acc:.4f}'),
'per_class_f1': per_class_dict,
'history': history
}
with open(os.path.join(LOGS_DIR, 'exp3_textrnn_results.json'), 'w') as f:
json.dump(result, f, ensure_ascii=False, indent=2)
log_csv = os.path.join(LOGS_DIR, 'experiment_log.csv')
entry = pd.DataFrame([{
'实验ID': 'exp3', '模型名称': 'TextRNN_BiLSTM_Attention', 'max_len': MAX_LEN,
'batch_size': BATCH_SIZE, 'lr': LR, 'epochs': EPOCHS,
'val_macro_f1': best_f1, 'val_accuracy': final_acc,
'训练时间(min)': '', '备注': f'hidden={HIDDEN_SIZE} layers={NUM_LAYERS}'
}])
entry.to_csv(log_csv, mode='a', header=not os.path.exists(log_csv), index=False)
print("\n[完成] 实验3 TextRNN完毕!")
if __name__ == '__main__':
main()
