pytorch - 为 iOS 导出时 Python 和 LibTorch C++ 之间的输出不一致

Question

我已经为我的数据训练了 HuggingFace RoBERTa 模型（这是一个非常特殊的用法——因此是小型模型/词汇表！）并在 Python 上成功测试。我将跟踪模型导出到 iOS 的 LibTorch，但设备上的预测结果与 Python 中的预测结果不匹配（给出不同的 argmax 令牌索引）。我的转换脚本：

# torch = 1.5.0
# transformers = 3.2.0

config = RobertaConfig(
    vocab_size=858,
    max_position_embeddings=258,
    num_attention_heads=6,
    num_hidden_layers=4,
    type_vocab_size=1,
    torchscript=True,
)

model = RobertaForMaskedLM(config=config).from_pretrained('./trained_RoBERTa')
model.cpu()
model.eval()

example_input = torch.LongTensor(1, 256).random_(0, 857).cpu()
traced_model = torch.jit.trace(model, example_input)
traced_model.save('./exports/trained_RoBERTa.pt')

过去，我在使用 Python+GPU 训练并转换为 iOS 的 LibTorch 的另一个（视觉）模型时遇到了问题，通过在我的转换脚本中添加调用map_location={'cuda:0': 'cpu'}来解决这些问题。所以我想知道：1）在这种情况下作为可能的解释是否有意义？2）在使用语法加载时torch.load()如何添加map_location选项？.from_pretrained()

以防万一我的 Obj-C++ 处理预测结果是罪魁祸首，这里是设备上运行的 Obj-C++ 代码：

- (NSArray<NSArray<NSNumber*>*>*)predictText:(NSArray<NSNumber*>*)tokenIDs {
    try {
        long count = tokenIDs.count;
        long* buffer = new long[count];
        for(int i=0; i < count;  i++) {
            buffer[i] = tokenIDs[i].intValue;
        }
        at::Tensor tensor = torch::from_blob(buffer, {1, (int64_t)count}, at::kLong);
        torch::autograd::AutoGradMode guard(false);
        at::AutoNonVariableTypeMode non_var_type_mode(true);
        auto outputTuple = _impl.forward({tensor}).toTuple();

        auto outputTensor = outputTuple->elements()[0].toTensor();
        auto sizes = outputTensor.sizes();
        // len will be tokens * vocab size -- sizes[1] * sizes[2] (sizes[0] is batch_size = 1)
        auto positions = sizes[1];
        auto tokens = sizes[2];
        float* floatBuffer = outputTensor.data_ptr<float>();
        if (!floatBuffer) {
            return nil;
        }
        // MARK: This is probably a slow way to create this 2D NSArray
        NSMutableArray* results = [[NSMutableArray alloc] initWithCapacity: positions];
        for (int i = 0; i < positions; i++) {
            NSMutableArray* weights = [[NSMutableArray alloc] initWithCapacity: tokens];
            for (int j = 0; j < tokens; j++) {
                [weights addObject:@(floatBuffer[i*positions + j])];
            }
            [results addObject:weights];
        }
        return [results copy];
    } catch (const std::exception& exception) {
        NSLog(@"%s", exception.what());
    }
    return nil;
}

请注意，我在 iOS 中的初始化代码确实调用eval()了 TorchScript 模型。

更新：一项观察；我在加载上面的训练模型时尝试使用我的方式config导致torchscript未设置标志 - 我认为它config完全忽略了我并从预训练文件中获取它。from_pretrained('./trained_RoBERTa', torchscript=True)因此，如文档中所述，我已将其移至。iOS上的输出也有同样的问题，请注意...

更新 2：我想我会尝试在 Python 中测试跟踪模型。不确定这是否应该起作用，但输出确实与原始模型中的相同测试匹配：

traced_test = traced_model(input)
pred = torch.argmax(traced_test[0], dim=2).squeeze(0)
pred_str = tokenizer.decode(pred[1:-1].tolist())
print(pred_str)

这让我觉得 iOS Obj-C++ 执行有问题。加载跟踪模型/导出的代码确实调用.eval()了模型，顺便说一句（我意识到这是对不同输出的可能解释）：

- (nullable instancetype)initWithFileAtPath:(NSString*)filePath {
    self = [super init];
    if (self) {
        try {
            auto qengines = at::globalContext().supportedQEngines();
            if (std::find(qengines.begin(), qengines.end(), at::QEngine::QNNPACK) != qengines.end()) {
                at::globalContext().setQEngine(at::QEngine::QNNPACK);
            }
            _impl = torch::jit::load(filePath.UTF8String);
            _impl.eval();
        } catch (const std::exception& exception) {
            NSLog(@"%s", exception.what());
            return nil;
        }
    }
    return self;
}

更新 3：Uhhhmmm ......这绝对是一个面对面的时刻（在浪费了一个周末之后）......我决定从 Obj-C 返回一个平面 NSArray 并在 Swift 中进行 2D 数组重塑，除了转变一个令牌（我认为它只是 [CLS]），输出现在是正确的。我猜我的 Obj-C 真的很生锈。可悲的是，我仍然没有看到这个问题，但它现在正在工作，所以我要投降了。

标签： pytorchobjective-c++libtorch