【模型转换】PyTorch模型转Caffe模型

最新推荐文章于 2024-06-29 19:20:56 发布

BILLY BILLY

最新推荐文章于 2024-06-29 19:20:56 发布

阅读量374

点赞数 1

分类专栏：剪枝蒸馏量化板端部署文章标签： pytorch caffe 人工智能

本文链接：https://blog.csdn.net/m0_51579041/article/details/136684435

版权

剪枝蒸馏量化板端部署专栏收录该内容

19 篇文章 6 订阅

订阅专栏

本文介绍了如何从GitHub上的PyTorchToCaffe库下载源码，并通过示例展示了如何将VGG11_bnPyTorch模型转换为Caffe模型，保存为prototxt和caffemodel文件。

摘要由CSDN通过智能技术生成

一、Caffe模型转换源码下载

GitHub:https://github.com/xxradon/PytorchToCaffe
GitHub上拉取PytorchToCaffe的源码，将Caffe文件夹和pytorch_to_caffe.py文件放到项目根目录
在这里插入图片描述

二、将VGG的PyTorch模型转为Caffe模型

import sys
sys.path.insert(0,'.')
import torch
from torch.autograd import Variable
from torchvision.models.vgg import vgg11_bn
import pytorch_to_caffe

if __name__=='__main__':
    name='vgg11_bn'
    net=vgg11_bn(True)
    input=Variable(torch.ones([1,3,224,224]))
    pytorch_to_caffe.trans_net(net,input,name)
    pytorch_to_caffe.save_prototxt('{}.prototxt'.format(name))
    pytorch_to_caffe.save_caffemodel('{}.caffemodel'.format(name))