在 GPU 上运行 BERT SQUAD model

Question

我正在使用 BERT Squad model 对一组文档 (>20,000) 提出相同的问题。 model 目前在我的 CPU 上运行，处理一个文档大约需要一分钟 - 这意味着我需要几天时间才能完成该程序。

我想知道是否可以通过在 GPU 上运行 model 来加快速度。 但是，我是 GPU 新手，我不知道如何将这些输入和 model 发送到设备 (Titan xp)。

该代码是从 Chris McChormick 那里借来的。

import torch
import tensorflow as tf
from transformers import BertForQuestionAnswering
from transformers import BertTokenizer

model = BertForQuestionAnswering.from_pretrained('bert-large-uncased-whole-word-masking-finetuned-squad')
tokenizer = BertTokenizer.from_pretrained('bert-large-uncased-whole-word-masking-finetuned-squad')

'question' 和 'answer_text' 分别是问题和上下文字符串。

    input_ids = tokenizer.encode(question, answer_text)

    # ======== Set Segment IDs ========
    # Search the input_ids for the first instance of the `[SEP]` token.
    sep_index = input_ids.index(tokenizer.sep_token_id)

    if len(input_ids)>512:
        input_ids=input_ids[:512]
  
    num_seg_a = sep_index + 1
    num_seg_b = len(input_ids) - num_seg_a

    # Construct the list of 0s and 1s.
    segment_ids = [0]*num_seg_a + [1]*num_seg_b

    # There should be a segment_id for every input token.
    assert len(segment_ids) == len(input_ids)

    # ======== Evaluate ========
    # Run our example through the model.
    outputs = model(torch.tensor([input_ids]), # The tokens representing our input text.
                    token_type_ids=torch.tensor([segment_ids]), # The segment IDs to differentiate question from answer_text
                    return_dict=True) 

    start_scores = outputs.start_logits
    end_scores = outputs.end_logits

我知道我可以使用 model.tocuda() 将 model 发送到 GPU。 但是我如何发送输入，训练 model，并从 GPU 中检索 output？

Answer 1

已经有一段时间了，但无论如何我都会回答，希望它可能会对某人有所帮助。 您可以使用to方法将每个张量复制到 GPU。 例如，您的批次包含 4 个 pytorch 张量：输入 ID、注意掩码、段 ID 和标签

device = torch.device("cuda")
b_input_ids = batch[0].to(device)
b_input_mask = batch[1].to(device)
b_seg_ids = batch[2].to(device)
b_labels = batch[2].to(device)

然后，您可以使用.cpu()将日志和标签从 gpu 传输回 cpu。 例如；

start_logits = start_logits.detach().cpu()
end_logits = end_logits.detach().cpu()

或与（设备）类似，您可以使用

start_logits = start_logits.to('cpu')
end_logits = end_logits.to('cpu')

请注意：由于您将在 model 中使用它们，您可能需要在末尾添加.numpy() 并将它们转换为 numpy 数组。

资料来源： https://discuss.pytorch.org/t/time-to-transform-gpu-to-cpu-with-cpu/18856