c++ - 从 c10::Dict 获取值在 Pytorch C++ 中

Question

我在 Pytorch C++ 前端使用 TorchScript 模型。

Python 中的模型将outputdict 返回为 Dict[str, List[torch.Tensor]].

当我在 C++ 中使用它时，它返回一个c10::Dict<c10::IValue, c10::IValue>. 这个 Python 代码的等价物是什么：

value_a = output['key_a']
value_b = output['key_b']

在 C++ 中获取价值c10::Dict？

我已经尝试过了，但它不起作用。

torch::IValue key_a("key_a");
torch::IValue key_b("key_b");
c10::IValue value_a = output[key_a];
c10::IValue value_b = output[key_b];
std::cout << value_a << std::endl;
std::cout << value_b << std::endl;

和错误：

error: type 'c10::Dict<c10::IValue, c10::IValue>' does not provide a subscript operator

Answer 1

c10:Dict 你可以在这里找到头文件。你想要的是at方法（在这里定义），所以：

auto value_a = output.at(key_a);

应该做的伎俩。

此外，您不必torch::IValue key_ay("key_a")显式创建，这应该足够了：

auto value_a = output.at("key_a");