caffe 模型转tensorflow

下载模型转换代码:

git clone https://github.com/ethereon/caffe-tensorflow

caffe-tensorflow模型转换文件为convert.py,可以将caffe的模型定义文件.prototxt和训练好的模型分别转换为tensorflow类,和tensorflow模型文件.在example文件夹下提供了examples/minist和examples/imagenet两个转换示例.

以手写字体为例,模型为lenet,examples/minist目录下提供了lenet.prototxt和训练好的caffe模lenet_iter_10000.caffemodel,caffe转换为tensorflow命令为:

python convert.py --def_path examples/mnist/lenet.prototxt --caffemodel examples/mnist/lenet_iter_10000.caffemodel --data-output-path lenet.npy --code-output-path lenet.py

运行命令后报错:

TypeError: Descriptors should not be created directly, but only retrieved from their parent.

Process finished with exit code 1

这是因为找不到caffe安装目录,解决方法是kaffe/caffe/resolver.py的第12行添加caffe安装目录:

import sys
caffe_root = '/home/qinghua/program/caffe'
sys.path.insert(0, caffe_root + '/python')
import caffe

转换后得到lenet网络结构文件lenet.py:

from kaffe.tensorflow import Network

class LeNet(Network):
    def setup(self):
        (self.feed('data')
             .conv(5, 5, 20, 1, 1, padding='VALID', relu=False, name='conv1')
             .max_pool(2, 2, 2, 2, name='pool1')
             .conv(5, 5, 50, 1, 1, padding='VALID', relu=False, name='conv2')
             .max_pool(2, 2, 2, 2, name='pool2')
             .fc(500, name='ip1')
             .fc(10, relu=False, name='ip2')
             .softmax(name='prob'))

已经模型文件lenet.npy,我们可以加载该模型,并进行fine tuning,代码为examples/mnist/finetune_mnist.py

# Import the converted model's class
import numpy as np
import random
import tensorflow as tf
from tensorflow.examples.tutorials.mnist import input_data

from lenet import LeNet as MyNet

mnist = input_data.read_data_sets('MNIST_data', one_hot=True)
batch_size = 32

def gen_data(source):
    while True:
        indices = range(len(source.images))
        random.shuffle(indices)
        for i in indices:
            image = np.reshape(source.images[i], (28, 28, 1))
            label = source.labels[i]
            yield image, label

def gen_data_batch(source):
    data_gen = gen_data(source)
    while True:
        image_batch = []
        label_batch = []
        for _ in range(batch_size):
            image, label = next(data_gen)
            image_batch.append(image)
            label_batch.append(label)
        yield np.array(image_batch), np.array(label_batch)


images = tf.placeholder(tf.float32, [batch_size, 28, 28, 1])
labels = tf.placeholder(tf.float32, [batch_size, 10])
net = MyNet({'data': images})

ip2 = net.layers['ip2']
pred = tf.nn.softmax(ip2)

loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=ip2, labels=labels), 0)
opt = tf.train.RMSPropOptimizer(0.001)
train_op = opt.minimize(loss)

with tf.Session() as sess:
    # Load the data
    sess.run(tf.initialize_all_variables())
    net.load('lenet.npy', sess)

    data_gen = gen_data_batch(mnist.train)
    for i in range(1000):
        np_images, np_labels = next(data_gen)
        feed = {images: np_images, labels: np_labels}

        np_loss, np_pred, _ = sess.run([loss, pred, train_op], feed_dict=feed)
        if i % 10 == 0:
            print('Iteration: ', i, np_loss)

更新prototxt及model为最新的caffe版本

由于caffe版本不同,一些参数定义会有更改,转换的时候由于不同caffe版本的问题,可能会出现错误:

Error encountered: Unknown layer type encountered: 4

这时候就需要使用函数upgrade_net_proto_text和函数upgrade_net_proto_binary,更新prototxt及model为最新的caffe版本.

更新prototxt:

/home/caffe/build/tools/upgrade_net_proto_text LightenedCNN_A_deploy.prototxt_old LightenedCNN_A_deploy.prototxt

更新model:

 /home/caffe/build/tools/upgrade_net_proto_binary LightenedCNN_A.caffemodel_old LightenedCNN_A.caffemodel

对于错误2:

Error encountered: Multiple top nodes are not supported.

这是因为prototxt文件中定义了包含两个top node的层,如:

layer {
  name: "slice1"
  type: "Slice"
  bottom: "pool1"
  top: "slice1_1"
  top: "slice1_2"
  slice_param {
    slice_dim: 1
  }
}

这时候就需要自己编写tensorflow代码.

但是完全自己编写也不太容易,这时候我们可以去掉不能转换的层,例如去掉Slice层,将Slice层的上下层连接,同时转换的时候只输入prototxt文件,即只对模型定义文件进行转换,以lightCNN为例,

python convert.py --def_path /home/face_verification_experiment/proto/LightenedCNN_A_deploy.prototxt  --code-output-path lightcnn.py

得到模型类为:

from kaffe.tensorflow import Network

class DeepFace_set003_net(Network):
    def setup(self):
        (self.feed('input')
             .conv(9, 9, 96, 1, 1, padding='VALID', relu=False, name='conv1')
             .max_pool(2, 2, 2, 2, name='pool1')
             .conv(5, 5, 192, 1, 1, padding='VALID', relu=False, name='conv2')
             .max_pool(2, 2, 2, 2, name='pool2')
             .conv(5, 5, 256, 1, 1, padding='VALID', relu=False, name='conv3')
             .max_pool(2, 2, 2, 2, name='pool3')
             .conv(4, 4, 384, 1, 1, padding='VALID', relu=False, name='conv4')
             .max_pool(2, 2, 2, 2, name='pool4')
             .fc(512, relu=False, name='fc1')
             .dropout(0.699999988079, name='dropout1')
             .fc(10575, relu=False, name='fc2')
             .softmax(name='softmax'))

参考以上代码,就可以使用tensorflow快速的编写对应的模型代码.

编写好模型代码后,需要加载训练好的caffe模型,我们加载caffe模型文件,并保存参数数组.

import sys

caffe_root = '/home/qinghua/program/caffe'
sys.path.insert(0, caffe_root + '/python')
# Try to import PyCaffe first
import caffe
import cPickle
dir='/home/face_verification_experiment/'
def_path=dir+'proto/LightenedCNN_A_deploy.prototxt_old'
data_path=dir+'model/LightenedCNN_A.caffemodel'
net = caffe.Net(def_path, data_path, caffe.TEST)
data = lambda blob: blob.data
params = [(k, map(data, v)) for k, v in net.params.items()]
params1=dict()
params1['params']=params
blobs=net.blobs
input_dims = blobs['data'].shape
with open(dir+'model/LightenedCNN_A.pkl', 'wb') as pickle_file:
    cPickle.dump(params1,pickle_file)

将所有训练好的变量值保存在LightenedCNN_A.pkl中,这样在可以使用tf.assign对tensorflow模型中的变量复制,再使用tf.saver.save()保存模型即可.这样最终完成了模型转换.