在 Pytorch C++ 中从 c10::Dict<c10::IValue, c10::IValue> 获取值

Question

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

Python returns 中的模型 output dict 为 Dict[str, List[torch.Tensor]]。

当我在 C++ 中使用它时，它 returns 一个 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:Dicthere. What you want is at method (defined here)的头文件，所以：

auto value_a = output.at(key_a);

应该可以解决问题。

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

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