怎样c++调用pytorch训练的模型

梦回兵工厂

已于 2022-05-17 10:19:41 修改

阅读量1.7k

点赞数 2

分类专栏： Deep Learning工程部署文章标签： pytorch 深度学习神经网络

于 2022-05-16 21:00:00 首次发布

本文链接：https://blog.csdn.net/m0_37635769/article/details/124795579

版权

Deep Learning工程部署专栏收录该内容

11 篇文章 0 订阅

订阅专栏

概括–常用思路：
思路1） pytorch框架模型转libtorch框架模型;
思路2）将pytorch下.pt模型先转通用的.onnx模型，再使用tensorrt加速工具转.engine模型（注：不同平台下的加速工具不同，例如Nvidia家tensorRT、Rockchip家RKNN）

一、思路1 pytorch环境模型转libtorch环境模型;
1、模型转换
首先在pytorch环境下，使用torch.jit.trace()或torch.jit.scrpit方法，生成libtorch环境需要的.pt模型。
下面以**torch.jit.trace()**为例：

if __name__ == '__main__':
    args = get_parser().parse_args()
    cfg = setup_cfg(args)

    cfg.defrost()
    cfg.MODEL.BACKBONE.PRETRAIN = False
    if cfg.MODEL.HEADS.POOL_LAYER == 'FastGlobalAvgPool':
        cfg.MODEL.HEADS.POOL_LAYER = 'GlobalAvgPool'
    model = build_model(cfg)  #！！重要
    Checkpointer(model).load(cfg.MODEL.WEIGHTS)  #！！重要
    if hasattr(model.backbone, 'deploy'):
        model.backbone.deploy(True)
    model.eval()

    inputs = torch.randn(args.batch_size, 3, cfg.INPUT.SIZE_TEST[0], cfg.INPUT.SIZE_TEST[1]).to(model.device)
    traced_script_module = torch.jit.trace(model, inputs) #！！重要
    traced_script_module.save("script.pt")

2、模型调用
使用相同或更高版本的libtorch，加载上一步骤生成的.pt模型。


    try {
        module_ = torch::jit::load(“xx.pt”);
    }
    catch (const c10::Error& e) {
        std::cerr << "Error loading the model!\n";
        std::exit(EXIT_FAILURE);
    }