常用的深度学习框架有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>