如何修改Python音频处理代码以用GPU替代RAM进行数据处理？

首先明确：音频特征提取和文本token化的基础操作多偏向CPU友好，但可以通过将生成的张量移至GPU，或结合PyTorch的GPU加速能力优化处理流程。以下是两种具体修改方案：

方案1：CPU预处理+整体移至GPU（推荐通用方案） #

先在CPU完成高效的多进程预处理，再将整个数据集移至GPU，避免GPU资源浪费在IO密集型的音频读取环节：

def prepare_dataset(batch):
    audio = batch["audio"]
    # 直接返回PyTorch格式的特征张量
    features = feature_extractor(
        audio["array"], 
        sampling_rate=audio["sampling_rate"],
        return_tensors="pt"
    ).input_features[0]
    # 生成标签张量并压缩多余维度
    labels = tokenizer(
        batch["sentence"],
        return_tensors="pt"
    ).input_ids.squeeze()
    return {"input_features": features, "labels": labels}

# 多进程加速CPU预处理（根据CPU核心数调整num_proc）
common_voice = common_voice.map(
    prepare_dataset, 
    remove_columns=common_voice.column_names["train"], 
    num_proc=4
)

# 将整个数据集格式转为torch并移至GPU
common_voice = common_voice.with_format("torch", device="cuda")

方案2：预处理阶段直接用GPU处理（适配支持GPU的特征提取器） #

如果你的feature_extractor（如WhisperFeatureExtractor）支持GPU张量输入，可在预处理时直接在GPU上操作：

import torch

def prepare_dataset(batch):
    audio = batch["audio"]
    # 将音频数组转为GPU张量
    audio_tensor = torch.tensor(audio["array"], device="cuda")
    # 在GPU上提取特征
    features = feature_extractor(
        audio_tensor, 
        sampling_rate=audio["sampling_rate"],
        return_tensors="pt"
    ).input_features[0].to("cuda")
    # token化后将标签移至GPU
    labels = tokenizer(
        batch["sentence"],
        return_tensors="pt"
    ).input_ids.squeeze().to("cuda")
    return {"input_features": features, "labels": labels}

# 多进程无法共享GPU上下文，需将num_proc设为1
common_voice = common_voice.map(
    prepare_dataset, 
    remove_columns=common_voice.column_names["train"], 
    num_proc=1
)

关键注意事项 #

• 并非所有特征提取器都支持GPU输入，需先确认你使用的feature_extractor兼容GPU张量。
• 方案1的效率通常更高：CPU负责IO和轻量预处理，GPU专注后续训练计算，避免GPU闲置。
• 多GPU环境下，可将device="cuda"替换为device="cuda:0"指定具体显卡，或借助accelerate库自动分配设备。

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