从跟踪的 torchscript 模型中获取图像大小答案

【问题标题】：Get image size from traced torchscript model从跟踪的 torchscript 模型中获取图像大小
【发布时间】：2023-01-07 05:48:09
【问题描述】：

将 yolo5 模型导出到 torchscript 时（使用 yolo5 的 export.py），必须使用 --imgsz 参数固定输入图像大小。

以后如何使用 torch 的 C++-API 访问此值（预期的图像大小）？

我知道信息在那里：导出的模型是伪装的 zip 存档，解压缩后，大小显示在 extra/config.txt 中。所以理论上我可以使用 zip 库来检索此信息。但是，它感觉很老套并且可能不稳定。

【问题讨论】：

标签： c++ torch yolo torchscript

【解决方案1】：

您可以使用 torch::jit::load 方法的重载获取 zip 存档中 extra/ 目录中的文件：

torch::jit::ExtraFilesMap extra_files{{"config.txt", ""}};
m_module = torch::jit::load("path/to/model.torchscript", gpu_id, extra_files);
std::cout << extra_files.at("config.txt") << std::endl;

将产生 config.txt 的内容，其中包含 JSON，即输入图像所需的形状。

从该字符串中获取所需的值应该很容易。

请注意，此解决方案仅在 config.txt 文件已明确添加到 zip 存档中时才有效。 yolov5's export.py does that。

【讨论】：

【解决方案2】：

我不是 C++ 专家，但这段代码对我有用。我在 python 上保存了两个 double 变量，并且能够在 C++ 中读取它们。我尝试直接从 string 到 memcpy 但我做不到，这就是我创建这个 char 序列的原因。

torch::jit::script::Module module;
torch::jit::ExtraFilesMap extra_files{{"var1", ""}, {"var2", ""}};

module = torch::jit::load("model.pt", device, extra_files);

std::string var1_str = extra_files.at("var1");
std::string var2_str = extra_files.at("var2");

double var1;
char b[] = {var1_str[0], var1_str[1], var1_str[2], var1_str[3], var1_str[4], var1_str[5], var1_str[6], var1_str[7]};
memcpy(&var1, &b, sizeof(var1));

printf("Var1 read: %.17g
", var1);

这是Python我用来保存变量的代码：

traced_script_module = torch.jit.trace(model, sample)

extra_files = {'var1': np.array(var1).tobytes(), 'var2': np.array(var2).tobytes()}

traced_script_module.save('model.pt', _extra_files=extra_files)

这是基于@pasbi 的回答。

【讨论】：