如何在Kaggle利用TPU加速预训练FinBert的财报电话会议情感分析？

核心结论 #

你必须在推理阶段将模型和数据整合到TPU上，默认的results = nlp(sentences)只会跑在CPU/GPU上，完全没用到TPU资源，这也是你当前耗时久的核心原因。

具体操作步骤 #

1. 初始化TPU环境 #

先确保正确加载TPU相关依赖并初始化设备：

import torch
import torch_xla.core.xla_model as xm
import torch_xla.distributed.xla_multiprocessing as xmp

# 获取TPU设备
device = xm.xla_device()

2. 手动加载FinBert模型到TPU #

Hugging Face的pipeline默认不支持TPU，需要手动加载tokenizer和模型并迁移到TPU：

from transformers import AutoTokenizer, AutoModelForSequenceClassification

# 加载FinBert的tokenizer和预训练模型
tokenizer = AutoTokenizer.from_pretrained("yiyanghkust/finbert-tone")
model = AutoModelForSequenceClassification.from_pretrained("yiyanghkust/finbert-tone")

# 将模型移到TPU设备，并切换到推理模式
model = model.to(device)
model.eval()

3. 修改推理循环，适配TPU #

将文本分批次处理，把每个批次的数据也移到TPU上进行推理，最后把结果移回CPU整理：

# 假设sentences是你的4万份财报电话会议文本列表
batch_size = 128  # 根据TPU内存调整，建议设大（比如256，只要不OOM）
all_results = []

for idx in range(0, len(sentences), batch_size):
    # 取出当前批次的文本
    batch_texts = sentences[idx:idx+batch_size]
    
    # 编码文本，生成TPU兼容的张量
    inputs = tokenizer(
        batch_texts,
        padding=True,
        truncation=True,
        max_length=512,
        return_tensors="pt"
    ).to(device)
    
    # 关闭梯度计算，加速推理
    with torch.no_grad():
        outputs = model(**inputs)
    
    # 将结果从TPU移回CPU，转换为numpy格式处理
    pred_ids = torch.argmax(outputs.logits, dim=-1).cpu().numpy()
    pred_scores = torch.softmax(outputs.logits, dim=-1).cpu().numpy()
    
    # 整理成和pipeline一致的结果格式
    for text, pred_id, scores in zip(batch_texts, pred_ids, pred_scores):
        label = model.config.id2label[pred_id]
        all_results.append({
            "text": text,
            "label": label,
            "score": float(scores[pred_id])
        })
    
    # TPU同步操作，确保当前批次计算完成（Kaggle环境下能稳定进度）
    xm.mark_step()

4. 进阶优化：多进程利用TPU全核心 #

Kaggle的TPU是8核心的，用多进程并行处理能最大化效率：

def process_inference(rank, all_sentences):
    # 每个进程绑定一个TPU核心
    device = xm.xla_device()
    tokenizer = AutoTokenizer.from_pretrained("yiyanghkust/finbert-tone")
    model = AutoModelForSequenceClassification.from_pretrained("yiyanghkust/finbert-tone").to(device).eval()
    
    # 将文本拆分到8个进程，每个进程处理1/8的数据
    process_sentences = all_sentences[rank::8]
    batch_size = 256
    process_results = []
    
    for idx in range(0, len(process_sentences), batch_size):
        batch_texts = process_sentences[idx:idx+batch_size]
        inputs = tokenizer(batch_texts, padding=True, truncation=True, return_tensors="pt").to(device)
        
        with torch.no_grad():
            outputs = model(**inputs)
        
        pred_ids = torch.argmax(outputs.logits, dim=-1).cpu().numpy()
        pred_scores = torch.softmax(outputs.logits, dim=-1).cpu().numpy()
        
        for text, pred_id, scores in zip(batch_texts, pred_ids, pred_scores):
            process_results.append({
                "text": text,
                "label": model.config.id2label[pred_id],
                "score": float(scores[pred_id])
            })
        
        xm.mark_step()
    
    # 保存当前进程的结果
    xm.save(process_results, f"tpu_results_rank_{rank}.pkl")

# 启动8个进程并行处理
xmp.spawn(process_inference, args=(sentences,), nprocs=8, start_method="fork")

# 合并所有进程的结果
import pickle
final_results = []
for rank in range(8):
    with open(f"tpu_results_rank_{rank}.pkl", "rb") as f:
        final_results.extend(pickle.load(f))

关键注意事项 #

• 不要用默认的pipeline：它没有TPU适配逻辑，必须手动加载模型并迁移设备
• 增大批次大小：TPU擅长并行处理大批次数据，批次越大，加速效果越明显（只要不超出TPU内存）
• 数据和模型必须同设备：所有输入张量必须移到TPU，否则会报错

内容的提问来源于stack exchange，提问作者Kyle_Stockton