# Copyright 2017 Chenxi Liu. All Rights Reserved. | |
# | |
# Licensed under the Apache License, Version 2.0 (the "License"); | |
# you may not use this file except in compliance with the License. | |
# You may obtain a copy of the License at | |
# | |
# http://www.apache.org/licenses/LICENSE-2.0 | |
# | |
# Unless required by applicable law or agreed to in writing, software | |
# distributed under the License is distributed on an "AS IS" BASIS, | |
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
# See the License for the specific language governing permissions and | |
# limitations under the License. | |
# ============================================================================== | |
# This code is modified from | |
# https://github.com/ppwwyyxx/tensorpack/blob/master/tensorpack/utils/loadcaffe.py | |
# sample usage: | |
# python caffemodel2npy.py deploy.prototxt | |
# ../deeplab/ResNet101/init.caffemodel ./model/ResNet101_init.npy | |
import numpy as np | |
import pdb | |
import re | |
class CaffeLayerProcessor(object): | |
def __init__(self, net): | |
self.net = net | |
self.layer_names = net._layer_names | |
self.param_dict = {} | |
self.processors = { | |
'Convolution': self.proc_conv, | |
'InnerProduct': self.proc_fc, | |
'BatchNorm': self.proc_bn, | |
'Scale': self.proc_scale | |
} | |
def process(self): | |
for idx, layer in enumerate(self.net.layers): | |
param = layer.blobs | |
name = self.layer_names[idx] | |
if 'res05' in name or 'res075' in name: | |
continue | |
if layer.type in self.processors: | |
name_ = self.rename(name) | |
dic = self.processors[layer.type](idx, name_, param) | |
self.param_dict.update(dic) | |
return self.param_dict | |
def rename(self, caffe_layer_name): | |
if caffe_layer_name.startswith('scale'): | |
caffe_layer_name = 'bn' + caffe_layer_name[5:] | |
NAME_MAP = {'bn_conv1': 'group_1/bn_conv1', | |
'conv1': 'group_1/conv1', | |
'fc1_voc12_c0': 'fc1_voc12/conv0', | |
'fc1_voc12_c1': 'fc1_voc12/conv1', | |
'fc1_voc12_c2': 'fc1_voc12/conv2', | |
'fc1_voc12_c3': 'fc1_voc12/conv3'} | |
if caffe_layer_name in NAME_MAP: | |
return NAME_MAP[caffe_layer_name] | |
s = re.search('([a-z]+)([0-9]+)([a-z]+)_', caffe_layer_name) | |
if s is None: | |
s = re.search('([a-z]+)([0-9]+)([a-z]+)([0-9]+)_', caffe_layer_name) | |
layer_block_part1 = s.group(3) | |
layer_block_part2 = s.group(4) | |
assert layer_block_part1 in ['a', 'b'] | |
layer_block = 0 if layer_block_part1 == 'a' else int(layer_block_part2) | |
else: | |
layer_block = ord(s.group(3)) - ord('a') | |
layer_type = s.group(1) | |
layer_group = s.group(2) | |
layer_branch = int(re.search('_branch([0-9])', caffe_layer_name).group(1)) | |
assert layer_branch in [1, 2] | |
if layer_branch == 2: | |
layer_id = re.search('_branch[0-9]([a-z])', caffe_layer_name).group(1) | |
layer_id = ord(layer_id) - ord('a') + 1 | |
else: | |
layer_id = 'add' | |
TYPE_DICT = {'res':'conv', 'bn':'bn'} | |
layer_type = TYPE_DICT[layer_type] | |
tf_name = 'group_{}_{}/block_{}/{}'.format( | |
int(layer_group), layer_block, layer_id, layer_type) | |
print caffe_layer_name, tf_name | |
return tf_name | |
def proc_conv(self, idx, name, param): | |
assert len(param) <= 2 | |
assert param[0].data.ndim == 4 | |
# caffe: ch_out, ch_in, h, w | |
W = param[0].data.transpose(2,3,1,0) | |
if len(param) == 1: | |
return {name + '/DW': W} | |
else: | |
return {name + '/DW': W, | |
name + '/biases': param[1].data} | |
def proc_fc(self, idx, name, param): | |
# TODO caffe has an 'transpose' option for fc/W | |
assert len(param) == 2 | |
prev_layer_name = self.net.bottom_names[name][0] | |
prev_layer_output = self.net.blobs[prev_layer_name].data | |
if prev_layer_output.ndim == 4: | |
W = param[0].data | |
# original: outx(CxHxW) | |
W = W.reshape((-1,) + prev_layer_output.shape[1:]).transpose(2,3,1,0) | |
# become: (HxWxC)xout | |
else: | |
W = param[0].data.transpose() | |
return {name + '/DW': W.squeeze(), | |
name + '/biases': param[1].data.squeeze()} | |
def proc_bn(self, idx, name, param): | |
# assert param[2].data[0] == 1.0 | |
return {name + '/mean': param[0].data, | |
name + '/variance': param[1].data, | |
name + '/factor': param[2].data } | |
def proc_scale(self, idx, name, param): | |
# bottom_name = self.net.bottom_names[name][0] | |
# # find the bn layer before this scaling | |
# for i, layer in enumerate(self.net.layers): | |
# if layer.type == 'BatchNorm': | |
# name2 = self.layer_names[i] | |
# bottom_name2 = self.net.bottom_names[name2][0] | |
# if bottom_name2 == bottom_name: | |
# # scaling and BN share the same bottom, should merge | |
# return {name2 + '/beta': param[1].data, | |
# name2 + '/gamma': param[0].data } | |
return {name + '/beta': param[1].data, | |
name + '/gamma': param[0].data} | |
# assume this scaling layer is part of some BN | |
# raise ValueError() | |
def load_caffe(model_desc, model_file): | |
""" | |
return a dict of params | |
""" | |
import caffe | |
caffe.set_mode_cpu() | |
net = caffe.Net(model_desc, model_file, caffe.TEST) | |
param_dict = CaffeLayerProcessor(net).process() | |
return param_dict | |
if __name__ == '__main__': | |
import argparse | |
parser = argparse.ArgumentParser() | |
parser.add_argument('model') | |
parser.add_argument('weights') | |
parser.add_argument('output') | |
args = parser.parse_args() | |
ret = load_caffe(args.model, args.weights) | |
# pdb.set_trace() | |
import numpy as np | |
np.save(args.output, ret) |
caffemodel2npy
最新推荐文章于 2021-04-12 01:02:58 发布