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

Posted 2023-03-29

【中文标题】从 Pytorch C++ 中的 c10::Dict<c10::IValue, c10::IValue> 获取值

【英文标题】：Get value from c10::Dict<c10::IValue, c10::IValue> in Pytorch C++

【发布时间】：2021-01-03 02:31:53

【问题描述】：

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

Python 中的模型返回 output 字典为 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

【问题讨论】：

C10:Dict 的 C++ 类型是什么？

@Surt。我不知道。 C10:Dict 来自 libtorch。你可以在这里阅读更多信息：pytorch.org/cppdocs/api/structc10_1_1_i_value.html

【参考方案1】：

你可以找到c10:Dicthere的头文件。你想要的是at方法（defined here），所以：

auto value_a = output.at(key_a);

应该做的伎俩。

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

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