PyTorch是否自动使用GPU？大文本语料嵌入GPU加速需求

你的代码速度慢的核心原因是模型和输入数据都运行在CPU上，即使CUDA可用也没利用到GPU算力。以下是具体修改方案：

关键修改步骤 #

• 加载模型时将其转移到GPU
• 将tokenizer生成的输入张量转移到GPU
• 支持批量处理文本（单条处理是海量数据慢的另一个关键因素）

修改后的完整代码 #

import torch
from transformers import BartTokenizer, BartForConditionalGeneration

model_name = 'facebook/bart-base'

def load_model():
    # 加载预训练模型并转移到GPU
    model = BartForConditionalGeneration.from_pretrained(model_name).to('cuda')
    # 加载tokenizer
    tokenizer = BartTokenizer.from_pretrained(model_name)
    return model, tokenizer

def calculate_text_embeddings(texts, model, tokenizer):
    # 批量tokenize文本，自动padding，返回张量转移到GPU
    inputs = tokenizer(
        texts, 
        return_tensors="pt", 
        max_length=512, 
        truncation=True,
        padding=True  # 开启批量padding
    ).to('cuda')
    
    with torch.no_grad():
        outputs = model.encoder(input_ids=inputs.input_ids, attention_mask=inputs.attention_mask)
    
    # 获取最后一层隐藏态，转移回CPU释放GPU显存
    last_hidden_states = outputs.encoder_last_hidden_state.cpu()
    
    return last_hidden_states

额外优化建议 #

• 批量处理: 调用函数时传入文本列表（而非单条文本），比如calculate_text_embeddings(["文本1", "文本2", ...], model, tokenizer)，批量大小根据你的GPU显存调整（例如64、128）
• 半精度加速: 如果GPU支持，可在加载模型时添加model = model.half()，用半精度模式进一步减少显存占用、提升计算速度
• 显存清理: 处理完一批数据后，调用torch.cuda.empty_cache()释放未使用的显存，避免显存溢出

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