美文网首页Python建模与NLP个人专题Pytorch与深度学习
C++部署pytorch模型(一)————将pytorch模型转

C++部署pytorch模型(一)————将pytorch模型转

作者: dalalaa | 来源:发表于2019-01-18 10:31 被阅读3次

    Pytorch官方提供的C++API名为libtorch,在Pytorch1.0版本之后开始支持windows,以后就可以直接用libtorch来部署Pytorch模型了。
    在Windows平台上用C++部署深度学习应用一直很麻烦,caffe使用自定义层还需要自己写backward(),而TensorFlow对Windows平台的支持也不是很好,在Windows上编译有各种坑等着你。而这次Facebook提供的libtorch还是编译好的文件,减少了很多的工作量,估计会吸引一部分人转向pytorch。

    下面是我将pytorch的模型转化为libtorch可用的torchscript模型的代码。

    import torch 
    import torchvision
    
     # 加载模型
    from models import FaceNetModel
    embedding_size = 128
    num_classes = 10000
    model = FaceNetModel(embedding_size = embedding_size,num_classes = num_classes)
    checkpoint = torch.load(r"D:\scripts\simple_facenet\log\checkpoint_epoch1610.pth")
    model.load_state_dict(checkpoint['state_dict'])
    model.cuda()
    model.eval()
    
    # 向模型中输入数据以得到模型参数 
    example = torch.rand(1,3,182,182).cuda() 
    traced_script_module = torch.jit.trace(model,example)
     
    # 保存模型
    traced_script_module.save("torch_script_eval.pt")
    

    这里有两个地方容易出错

    1. 如果模型中有dropout或者batchnorm的话,一定要先将模型设置为eval模式,再保存,否则在用libtorch调用后会出现随机干扰;
    2. example这个张量的尺寸务必与你自己的模型的输入尺寸一直,否则会出现错误。

    这个libtorch出来时间不长,还有很多坑等着我去踩,如果有同样在研究pytorch或者libtorch的朋友,可以加个好友,共同学习。

    相关文章

      网友评论

        本文标题:C++部署pytorch模型(一)————将pytorch模型转

        本文链接:https://www.haomeiwen.com/subject/swhudqtx.html