Search code examples
c++torchyolotorchscript

Get image size from traced torchscript model


When exporting a yolo5-model to torchscript (using yolo5's export.py), the input image size has to be fixed using the --imgsz-argument.

How can I access this value (the expected image size) later using torch's C++-API?

I know that the information is there: the exported model is a zip-archive in disguise and, after unzipping, the size appears in extra/config.txt. So I could theoretically use a zip-library to retrieve this information. However, it feels hacky and might be unstable.


Solution

  • You can get the files in the extra/ directory in the zip archive using an overload of the torch::jit::load method:

    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;
    

    Will yield the content of config.txt, which contains JSON, i.a. the required shape of the input image.

    It should be easy to troll the desired value from that string.

    Note that this solution only works if the config.txt file has been added to the zip archive explicitly. yolov5's export.py does that.