教你自己训练的pytorch模型转caffe(二)

重点之重点

把训练的模型进行转换,我的模型文件叫google_checkpoint_ep60

参考:https://github.com/inisis/brocolli
https://www.bilibili.com/video/BV1fb4y177Ry

那个老哥给我们搭建好了环境, 因此不用自己搭建环境,但是我们得下载docker

看这个视频把docker 装好 https://www.bilibili.com/video/BV11L411g7U1

然后我们就下载老哥提供的镜像

具体的,在命令行按照老哥的github指引敲

教你自己训练的pytorch模型转caffe(二)_第1张图片

 第一行是下载镜像(可以理解为虚拟机中的文件),第二行是创建容器(创建一个虚拟机然后开机了,并登录),我推荐大家把--rm去了,不然cmd退出后容器就自动删除了,第三行是进入容器执行命令

这个命令是转换模型,最后保存在tmp文件夹中。

教你自己训练的pytorch模型转caffe(二)_第2张图片

 如果没有报错的话下面就开始转换我们自己的模型

我使用docker cp 在宿主机和容器之间复制

首先把我们猫狗的GoogLeNet.py 复制进去。

docker cp E:/workspace/brocolli/custom_models/GoogLeNet.py ee132b3c68f8b98f204de3bbb8872573e9fb00db319a3dd8bb2c3674f6d4a776:/root/brocolli/custom_models 

前面是宿主机文件或者文件夹,后面是容器id ,id可以从docker desktop复制

教你自己训练的pytorch模型转caffe(二)_第3张图片教你自己训练的pytorch模型转caffe(二)_第4张图片

 最后是容器的文件夹

然后我们写一个脚本转换我们的模型,把老哥readme抄抄

教你自己训练的pytorch模型转caffe(二)_第5张图片

 我的文件名叫run.py

import torchvision
import torch

from custom_models import GoogLeNet
model = GoogLeNet.GoogLeNet(num_classes=2) # Here, you should use your ownd model
model.load_state_dict(torch.load("./google_checkpoint_ep60.pth",map_location=torch.device('cpu')),)

# input  = torch.rand(1,3,224,224)
# output = model(input)
# aaa = 100

from bin.pytorch2caffe import Runner # if caffe, use bin.pytorch2caffe, if TensorRT use bin.pytorch2trt;
runner = Runner("googlenet_dog", model, [1, 3, 224, 224], 13,True)

runner.pyotrch_inference()
runner.convert()
runner.caffe_inference() # if caffe, use caffe_inference, if TensorRT use trt_inference;
runner.check_result()

我们把run.py 和google_checkpoint_ep60.pth 都用docker cp 复制到/root/brocolli 中

然后运行,转换完毕,把 tmp下面的googlenet_dog.prototxt 和googlenet_dog.caffemodel 复制走至此完毕

分割线---------------------------------

docker 一些命令

docker run 创建容器

docker run -it  yaphets4desmond/brocolli:v1.0 bash
如果你之前用了rm 删掉了,用上述命令在创建一个,或者这样

 如果你小心把bash 窗口关了,但是他还在运行

 可以使用命令进入bash,长长的那个是容器id

docker exec -it  e132bd2faff7adb3596068d0f53316fc54d20309f2aed6464984206f3a811fe4 bash 

你可能感兴趣的:(海思AI,caffe,pytorch,深度学习)