caffemodel2npy

  # 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)
  • 0
    点赞
  • 0
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值