使用c++调用pytorch模型

 常用的深度学习框架有pytorch、tensorflow和caffe。caffe配置太麻烦，开发维护不方便。tensorflow、pytorch都不错，开发维护都很方便。

    目前深度学习大多用python进行开发维护，但某些场景需要用c++。以pytorch为例，使用c++进行开发，可以大幅提高运行速度。

    首先要把模型导出：

#coding=utf-8

import torch

from ** import models#引入定义好的模型

 

model=models()

pth_path=r"**.pt" #引入训练好的参数

checkpoint=torch.load(pth_path) #加载参数

model.load_state_dict(checkpoint)

model.cuda() #模型在GPU上运行

model.eval() #模型用来测试，而非训练

 

example1=torch.rand(1, 3, 512, 512).cuda() #生成测试sample

with torch.no_grad(): #不要梯度，不然显存会爆炸

    trace_script_modile1=torch.jit.trace(model, example1)

    trace_script_modile1.save(r"**.pt") #压缩好的模型存出来

 

c++调用：引入头文件和lib即可

#include <iostream>

#include <memory>

#include <ctime>