Pytorch_模型轉Caffe(三)pytorch轉caffemodel
阿新 • • 發佈:2020-12-21
目錄
3.
5.
Pytorch_模型轉Caffe(三)pytorch轉caffemodel
- 模型轉換基於GitHub上xxradon的程式碼進行優化,在此對作者表示感謝。GitHub地址:https://github.com/xxradon/PytorchToCaffe
- 本文基於AlexNet網路對MNIST手寫字型分類生成的模型*.pth進行轉換
1. Pytorch下生成模型
- 呼叫
torchvision.models.alexnet
下的alexnet
網路 - 修改網路輸入層數 1 ,輸出類別數量 10
classifier
下的dropout
位置需要調整
- 通過一下程式碼訓練手寫數字識別,最終生成模型
mnist_alexnet_model.pth
(這裡儲存了整個網路和權重)
import time import torch from torch import nn, optim import torchvision import pytorch_deep as pyd from torchvision.models.alexnet import alexnet net = alexnet(False) device = torch.device('cuda' if torch.cuda.is_available() else'cpu') def load_data_fashion_mnist(batch_size = 256,resize=None,num_workers = 0): trans = [] if resize: trans.append(torchvision.transforms.Resize(size=resize)) trans.append(torchvision.transforms.ToTensor()) transform = torchvision.transforms.Compose(trans) mnist_train = torchvision.datasets.FashionMNIST(root='./MNIST', train=True, download=True, transform=transform) mnist_test = torchvision.datasets.FashionMNIST(root='./MNIST', train=False, download=True, transform=transform) train_iter = torch.utils.data.DataLoader(mnist_train, batch_size=batch_size, shuffle=True, num_workers=num_workers) test_iter = torch.utils.data.DataLoader(mnist_test, batch_size=batch_size, shuffle=False, num_workers=num_workers) return train_iter,test_iter batch_size = 128 # 如出現“out of memory”的報錯資訊,可減⼩batch_size或resize train_iter, test_iter = load_data_fashion_mnist(batch_size,resize=224) lr, num_epochs = 0.001, 5 optimizer = torch.optim.Adam(net.parameters(), lr=lr) pyd.train_ch5(net, train_iter, test_iter, batch_size, optimizer,device, num_epochs)
2. pth轉換成caffemodel和prototxt
git clone
下載GitHub原始碼,進入example
下的Alexnet例項- 主要用到以下兩個檔案,一個是載入網路模型,一個是進行prototxt和caffemodel的轉換
- 先看
alexnet_pytorch_to_caffe.py
import sys sys.path.insert(0,'.') import torch from torch.autograd import Variable from torchvision.models.alexnet import alexnet import pytorch_to_caffe_alexNet import cv2 if __name__=='__main__': name='alexnet' pth_path = '***/mnist_alexnet_model.pth' net = torch.load(pth_path) net.eval() input=Variable(torch.FloatTensor(torch.ones([1,1,224,224]))) input = input.cuda() pytorch_to_caffe_alexNet.save_prototxt('{}.prototxt'.format(name)) pytorch_to_caffe_alexNet.save_caffemodel('{}.caffemodel'.format(name))
- 如果直接執行發現會報錯,我這裡的錯誤出現在
dropout
層轉化的位置,修改其bottom
和top
傳參 - 修改完
dropout
,執行正常,能夠生產caffemodel和prototxt,但prototxt網路結構有問題,還是前後層銜接不對 - 參照原版deploy.prototxt進行layer的修改,最終輸出了正確的結果
3. pytorch_to_caffe_alexNet.py
剖析
- 該檔案就是對pth檔案進行解析,獲得layer的名稱和每層的權重偏差,並以caffe的格式進行儲存
- 修改了pytorch Function中的函式,讓其在前向傳播的時候自動將該層的引數儲存到caffe
- 很多層的前後銜接不對,都需要強制進行修改
- 下面是修改的部分函式
def _dropout(raw,input,p=0.5, training=False, inplace=False):
x=raw(input,p, training, False)
layer_name=log.add_layer(name='dropout')
log.add_blobs([x],name='dropout_blob')
bottom_top_name = 'fc_blob' + layer_name[-1]
layer=caffe_net.Layer_param(name=layer_name,type='Dropout',
bottom=[bottom_top_name],top=[bottom_top_name])
layer.param.dropout_param.dropout_ratio = p
log.cnet.add_layer(layer)
return x
def _linear(raw,input, weight, bias=None):
x=raw(input,weight,bias)
layer_name=log.add_layer(name='fc')
top_blobs=log.add_blobs([x],name='fc_blob')
bottom_name = 'ave_pool_blob1' if top_blobs[-1][-1] =='1' else 'fc_blob'+str(int(top_blobs[-1][-1])-1)
layer=caffe_net.Layer_param(name=layer_name,type='InnerProduct',
bottom=[bottom_name],top=top_blobs)
layer.fc_param(x.size()[1],has_bias=bias is not None)
if bias is not None:
layer.add_data(weight.cpu().data.numpy(),bias.cpu().data.numpy())
else:
layer.add_data(weight.cpu().data.numpy())
log.cnet.add_layer(layer)
return x
4. 用轉換後的模型進行推理
- 在caffe 下進行測試 test_alexnet.sh
#!/bin/bash
set -e
./build/examples/cpp_classification/classification.bin \
/home/****/alexnet.prototxt \
/home/****/alexnet.caffemodel \
examples/mnist/mnist_mean.binaryproto \
examples/mnist/mnist_label.txt \
data/mnist/1.png;
目前推理結果不太準,但整個過程都已經跑通
5. prototxt
注意問題
-
推理過程發現每次的結果都不一樣,發現prototxt中每個卷積層下都有初始化權重的偏差,將其統統刪除
-
池化層下的
ceil_mode: false
也是多餘項,刪除即可
至此已完成Pytorch到caffemodle的轉換
這只是初步嘗試通過,接下來要進行YOLOv4的轉換,應該會遇到更多的問題,加油!