Skip to content

Commit

Permalink
simplify code
Browse files Browse the repository at this point in the history
  • Loading branch information
@ppwwyyxx
ppwwyyxx committed Oct 23, 2016
1 parent fc81be3 commit fcdeafb
Showing 1 changed file with 60 additions and 93 deletions.
153 changes: 60 additions & 93 deletions examples/ResNet/load-resnet.py
100644 → 100755
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,10 +10,12 @@
import argparse
import os, re
import numpy as np
import six
from six.moves import zip
from tensorflow.contrib.layers import variance_scaling_initializer

from tensorpack import *
from tensorpack.utils import logger
from tensorpack.tfutils.symbolic_functions import *
from tensorpack.tfutils.summary import *
from tensorpack.dataflow.dataset import ILSVRCMeta
Expand All @@ -32,14 +34,14 @@ def _get_input_vars(self):
def _build_graph(self, input_vars):
image = input_vars[0]

def caffe_shortcut(l, n_in, n_out, stride):
def shortcut(l, n_in, n_out, stride):
if n_in != n_out:
l = Conv2D('convshortcut', l, n_out, 1, stride=stride)
return BatchNorm('bnshortcut', l)
else:
return l

def caffe_bottleneck(l, ch_out, stride, preact):
def bottleneck(l, ch_out, stride, preact):
ch_in = l.get_shape().as_list()[-1]
input = l
if preact == 'both_preact':
Expand All @@ -53,143 +55,96 @@ def caffe_bottleneck(l, ch_out, stride, preact):
l = tf.nn.relu(l)
l = Conv2D('conv3', l, ch_out * 4, 1)
l = BatchNorm('bn3', l) # put bn at the bottom
return l + caffe_shortcut(input, ch_in, ch_out * 4, stride)
return l + shortcut(input, ch_in, ch_out * 4, stride)

def layer(l, layername, block_func, features, count, stride, first=False):
def layer(l, layername, features, count, stride, first=False):
with tf.variable_scope(layername):
with tf.variable_scope('block0'):
l = block_func(l, features, stride,
l = bottleneck(l, features, stride,
'no_preact' if first else 'both_preact')
for i in range(1, count):
with tf.variable_scope('block{}'.format(i)):
l = block_func(l, features, 1, 'both_preact')
l = bottleneck(l, features, 1, 'both_preact')
return l

cfg = {
50: ([3,4,6,3], caffe_bottleneck),
101: ([3,4,23,3], caffe_bottleneck),
152: ([3,8,36,3], caffe_bottleneck)
50: ([3,4,6,3]),
101: ([3,4,23,3]),
152: ([3,8,36,3])
}

defs, block_func = cfg[MODEL_DEPTH]

defs = cfg[MODEL_DEPTH]
with argscope(Conv2D, nl=tf.identity, use_bias=False,
W_init=variance_scaling_initializer(mode='FAN_OUT')):
fc1000l = (LinearWrap(image)
.Conv2D('conv0', 64, 7, stride=2, nl=BNReLU )
fc1000 = (LinearWrap(image)
.Conv2D('conv0', 64, 7, stride=2, nl=BNReLU)
.MaxPooling('pool0', shape=3, stride=2, padding='SAME')
.apply(layer, 'group0', block_func, 64, defs[0], 1, first=True)
.apply(layer, 'group1', block_func, 128, defs[1], 2)
.apply(layer, 'group2', block_func, 256, defs[2], 2)
.apply(layer, 'group3', block_func, 512, defs[3], 2)
.apply(layer, 'group0', 64, defs[0], 1, first=True)
.apply(layer, 'group1', 128, defs[1], 2)
.apply(layer, 'group2', 256, defs[2], 2)
.apply(layer, 'group3', 512, defs[3], 2)
.tf.nn.relu()
.GlobalAvgPooling('gap')
.FullyConnected('fc1000', 1000, nl=tf.identity)())
prob = tf.nn.softmax(fc1000, name='prob_output')

prob = tf.nn.softmax(fc1000l, name='prob_output')

def run_test(path, input):
def run_test(params, input):
image_mean = np.array([0.485, 0.456, 0.406], dtype='float32')
param = np.load(path).item()

resNet_param = { caffeResNet2tensorpackResNet(k) :v for k, v in param.iteritems()}
pred_config = PredictConfig(
model=Model(),
input_var_names=['input'],
session_init=ParamRestore(resNet_param),
session_init=ParamRestore(params),
output_var_names=['prob_output']
)
predict_func = get_predict_func(pred_config)

remap_func = lambda x: (x - image_mean * 255)
im = cv2.imread(input)
im = remap_func(cv2.resize(im, (224,224)))
im = cv2.resize(im, (224,224)) - image_mean * 255
im = np.reshape( im, (1, 224, 224, 3)).astype('float32')
input = [im]
prob = predict_func(input)[0]
prob = predict_func([im])[0]

ret = prob[0].argsort()[-10:][::-1]
print(ret)

meta = ILSVRCMeta().get_synset_words_1000()
print([meta[k] for k in ret])


def caffeResNet2tensorpackResNet(caffe_layer_name):
# begining & ending stage
name_map = {'bn_conv1/beta': 'conv0/bn/beta',
def name_conversion(caffe_layer_name):
# beginning & end mapping
NAME_MAP = {'bn_conv1/beta': 'conv0/bn/beta',
'bn_conv1/gamma': 'conv0/bn/gamma',
'bn_conv1/mean/EMA': 'conv0/bn/mean/EMA',
'bn_conv1/variance/EMA': 'conv0/bn/variance/EMA',
'conv1/W': 'conv0/W',
'conv1/b': 'conv0/b',
'fc1000/W': 'fc1000/W',
'fc1000/b': 'fc1000/b'}
if caffe_layer_ in name_map:
print(caffe_layer_name + ' --> ' + name_map[caffe_layer_name])
return name_map[caffe_layer_name]

print(caffe_layer_name)

layer_id = None
layer_type = None
layer_block = None
layer_branch = None
layer_group = None
s = re.search('([a-z]*)([0-9]*)([a-z]*)_branch([0-9])([a-z])', caffe_layer_name, re.IGNORECASE)
if s == None:
s = re.search('([a-z]*)([0-9]*)([a-z]*)_branch([0-9])', caffe_layer_name, re.IGNORECASE)
else:
layer_id = s.group(5)
if caffe_layer_name in NAME_MAP:
return NAME_MAP[caffe_layer_name]

if s.group(0) == caffe_layer_name[0:caffe_layer_name.index('/')]:
layer_type = s.group(1)
layer_group = s.group(2)
layer_block = ord(s.group(3)) - ord('a')
layer_branch = s.group(4)
else:
# print('s group ' + s.group(0))
s = re.search('([a-z]*)([0-9]*)([a-z]*)([0-9]*)_branch([0-9])([a-z])', caffe_layer_name, re.IGNORECASE)
if s == None:
s = re.search('([a-z]*)([0-9]*)([a-z]*)([0-9]*)_branch([0-9])', caffe_layer_name, re.IGNORECASE)
else:
layer_id = s.group(6)

layer_type = s.group(1)
layer_group = s.group(2)
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)
if layer_block_part1 == 'a':
layer_block = 0
elif layer_block_part1 == 'b':
layer_block = int(layer_block_part2)
else:
print('model block error!')

layer_branch = s.group(5)

if s.group(0) != caffe_layer_name[0:caffe_layer_name.index('/')]:
print('model depth error!')
# TODO error handling
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

type_dict = {'res': '/conv', 'bn':'/bn', 'scale':'/bn'}
shortcut_dict = {'res': '/convshortcut', 'bn':'/bnshortcut', 'scale':'/bnshortcut'}
TYPE_DICT = {'res':'conv', 'bn':'bn'}

tf_name = caffe_layer_name[caffe_layer_name.index('/'):]

if layer_branch == '2':
tf_name = 'group' + str( int(layer_group) - int('2') ) + \
'/block' + str( layer_block ) + \
type_dict[layer_type] + str( ord(layer_id) - ord('a') + 1) + tf_name
elif layer_branch == '1':
tf_name = 'group' + str( int(layer_group) - int('2') ) + \
'/block' + str(layer_block) + \
shortcut_dict[layer_type] + tf_name
else:
print('renaming error!')
# TODO error handling
print(caffe_layer_name + ' --> ' + tf_name)
layer_type = TYPE_DICT[layer_type] + \
(str(layer_id) if layer_branch == 2 else 'shortcut')
tf_name = 'group{}/block{}/{}'.format(
int(layer_group) - 2, layer_block, layer_type) + tf_name
return tf_name

if __name__ == '__main__':
Expand All @@ -206,4 +161,16 @@ def caffeResNet2tensorpackResNet(caffe_layer_name):
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
# run resNet with given model (in npy format)
MODEL_DEPTH = args.depth
run_test(args.load, args.input)

param = np.load(args.load, encoding='latin1').item()
resnet_param = {}
for k, v in six.iteritems(param):
try:
newname = name_conversion(k)
except:
logger.error("Exception when processing caffe layer {}".format(k))
raise
logger.info("Name Transform: " + k + ' --> ' + newname)
resnet_param[newname] = v

run_test(resnet_param, args.input)

0 comments on commit fcdeafb

Please sign in to comment.