c++ - 从Pytorch C++中的c10::Dict <c10::IValue，c10::IValue>获取值

Question

我在Pytorch C++前端上使用了TorchScript模型。
Python中的模型将output dict返回为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获得值(value)？
我已经尝试过了，但是没有用。

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 here的头文件。您想要的是at方法(defined here)，因此:

auto value_a = output.at(key_a);

应该做到的。
另外，您不必显式创建torch::IValue key_ay("key_a")，这应该足够了:

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