对大数据进行预测的最有效方法？

Question

我正在尝试编写一个预测函数来预测大型文本数据（所以它必须是批量的）。但是预测功能有点慢。所以我想知道我能做些什么来改善它的时间。

到目前为止我所拥有的：

def get_embeddings(model, data_loader, device):
    
    model.eval()  # eval mode
    
    with torch.no_grad():
        embeddings = torch.tensor([], dtype=torch.float64, device=device)  #initialize empty tensor
        for _, d in enumerate(tqdm(data_loader), 0):
            # model inputs
            input_ids = d["input_ids"].to(device)
            attention_mask = d["attention_mask"].to(device)

            # model outputs
            embeddings = torch.cat((embeddings, model.predict(
                input_ids,
                attention_mask=attention_mask
            )))  # concat predictions

    return embeddings.cpu().numpy()  # convert to numpy array

我认为从 GPU 转换到 CPU 需要时间，所以我决定先初始化一个空张量并连接所有预测。然后在最后将其转换回 numpy 数组。但我不确定张量连接是否真的会更慢。

所以我想知道在预测方面是否有更好或最佳实践。

Answer 1

您的代码中有几件事引起了我的注意：

• 您正在使用float64。 GPU 比 float32 和 float16 慢得多。
• 您可能会强制将 float32 转换为 float64，这也可能很慢。
• 您正在为每个输入将注意力掩码复制到 GPU。相反，创建和/或保留它。
• 确保所有操作均分批完成，以充分利用 GPU 的并行性。
• 连接可能没用：GPU->CPU 传输大多受带宽限制，因此它是一个大张量还是许多小张量（在一定程度上）并不重要。其次，分配（由cat 引起）可能很慢。