如何在 GPU 而不是 RAM 上处理此 Python 代码的数据？

Question

我当前正在使用以下代码来处理音频数据，但它在 RAM 上运行。我想将处理任务转移到 GPU 以提高性能。 我的代码：

def prepare_dataset(batch):
    audio = batch["audio"]
    batch["input_features"] = feature_extractor(
        audio["array"], 
        sampling_rate=audio["sampling_rate"]
    ).input_features[0]
    batch["labels"] = tokenizer(batch["sentence"]).input_ids
    return batch

common_voice = common_voice.map(
    prepare_dataset, 
    remove_columns=common_voice.column_names["train"], 
    num_proc=1
)

如何修改此代码以利用 GPU 而不是 RAM 进行处理？非常感谢任何指导或具体更改！

Answer 1

您可以使用以下代码在GPU上处理音频数据

import torch
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(device)

def prepare_dataset(batch):
    audio = batch["audio"]

    input_features = feature_extractor(audio["array"], sampling_rate=audio["sampling_rate"]).input_features[0]
    batch["input_features"] = torch.tensor(input_features).to(device)

    labels = tokenizer(batch["sentence"]).input_ids
    batch["labels"] = torch.tensor(labels).to(device)
    return batch

common_voice = common_voice.map(prepare_dataset, remove_columns=common_voice.column_names["train"])