caffe是比较老的框架了,pytorch还不火的时候,还是比较流行的,有些比较著名的如人脸识别网络如centerloss,目标检测网络mtcnn、ssd,OCR识别都有对应的caffe版本。但有几个问题:1、添加新的层比较麻烦,要写反向传播;2、搭建网络时,prototxt的网络结构比较麻烦,动辄几千行。3、一些新的trick添加比较麻烦。
pytorch可以解决以上问题。现在很多任务都会使用pytorch进行训练了。但pytorch也存在问题,在服务器部署时还好。但要是在嵌入式端部署,一般都需要把模型转为ncnn之类的arm推理框架支持的格式,或者转为带npu算力的芯片各自推理工具链转换的格式(如海思的nnie,瑞芯微的rknn,Ambarella的cv28等等)。问题来了,从海思开始,就是对caffe的支持比较好,其他几家的芯片也都慢慢开始支持pytorch等框架,但在部署时,得到的实验结果总是caffe量化完后精度损失相对更小,更为稳定。于是在后来的模型部署时,啥模型都转为caffe,工具链转换一把梭,基本没啥问题。所以,需要将pytorch训练的模型转换为caffe版本。
具体的模型转换有两种方式,一般一些简单的模型,可以直接使用github的工程
https://github.com/xxradon/PytorchToCaffe
,进行转换。我在pytorch1.2,torchvision0.2.2的环境中测试过resnet18和自己训练的小的分类模型,没啥大问题。但注意,caffe和pytorch在maxpooling的时候是ceil模式,还是floor模式,保持一致即可。
但有些目标检测、gan或者其他自己设计的使用了些复杂算子的模型,建议不这么做。主要是为了保证pytorch转为caffe无损转换,所有推理结果的tensor在caffe和pytorch下一致,还是手动操作更为靠谱,主要过程可控,也方便一层层调试。
该过程主要有以下几步:
1、根据pytorch的网络结构,手写一个caffe的网络结构net.prototxt,和训练solver.prototxt。再根据prototxt,生成对应的随机初始化的caffemodel。
import caffe
solver_file='model_solver.prototxt'
solver=caffe.get_solver(solver_file)
solver.net.save('random_init.cafffmodel')
2、将pytorch的保存下来的训练模型的权值转为numpy格式
import torch
model_file="pytorch_model.pth.tar'
state_dict=torch.load(model_file,map_location=lamda storage,loc:storage)
state_dict=state_dict['state_dict']
for param_tensor in state_dict():
print(param_tensor ,"\t",state_dict()[param_tensor].size())
for key in state_dict:
state_dict[key]=state_dict[key].numpy()
np.save("pytorch_model.npy",state_dict)
3、将pytorch的.npy格式权重转为caffe的caffemodel
import caffe
net_file="caffe_model.prototxt"
model_file="caffe_model.caffemodel"
net=caffe.Net(net_file,model_file,caffe.Test)
params=net.params
import numpy as np
pytorch_model=np.load("pytorch_model.npy")
pytorch_model=pytorch_model.tolist()
params['conv1'][0].data[...]=pytorch_model['conv1.weight'].cpoy()
params['conv1'][1].data[...]=pytorch_model['conv1.bias'].cpoy()
params['scal1'][0].data[...]=pytorch_model['bn1.weight'].cpoy()
params['scal1'][1].data[...]=pytorch_model['bn1.bias'].cpoy()
params['bn1'][2].data[...]=1
params['bn1'][0].data[...]=pytorch_model['bn1.running_mean'].cpoy()
params['bn1'][1].data[...]=pytorch_model['bn1.running_var'].cpoy()
params['fc1'][0].data[...]=pytorch_model['fc1.weight'].cpoy()
params['fc1'][1].data[...]=pytorch_model['fc1.bias'].cpoy()
#省略类似步骤
net.save("new_write.caffemodel")
后续就是自己验证比较caffe和pytorch的推理结果是否一致了,可逐层比较。
当然,以上步骤还是比较麻烦,后续还是期待各家的芯片对pytorch的支持也可以越来越好。