mirror of
https://github.com/wassname/PSPNet-Keras-tensorflow.git
synced 2026-09-11 11:50:41 +08:00
Made commit of original caffe-tensorflow converter
This commit is contained in:
VladK
@@ -0,0 +1,2 @@
|
||||
from .transformer import TensorFlowTransformer
|
||||
from .network import Network
|
||||
@@ -0,0 +1,244 @@
|
||||
import numpy as np
|
||||
import tensorflow as tf
|
||||
|
||||
DEFAULT_PADDING = 'SAME'
|
||||
|
||||
|
||||
def layer(op):
|
||||
'''Decorator for composable network layers.'''
|
||||
|
||||
def layer_decorated(self, *args, **kwargs):
|
||||
# Automatically set a name if not provided.
|
||||
name = kwargs.setdefault('name', self.get_unique_name(op.__name__))
|
||||
# Figure out the layer inputs.
|
||||
if len(self.terminals) == 0:
|
||||
raise RuntimeError('No input variables found for layer %s.' % name)
|
||||
elif len(self.terminals) == 1:
|
||||
layer_input = self.terminals[0]
|
||||
else:
|
||||
layer_input = list(self.terminals)
|
||||
# Perform the operation and get the output.
|
||||
layer_output = op(self, layer_input, *args, **kwargs)
|
||||
# Add to layer LUT.
|
||||
self.layers[name] = layer_output
|
||||
# This output is now the input for the next layer.
|
||||
self.feed(layer_output)
|
||||
# Return self for chained calls.
|
||||
return self
|
||||
|
||||
return layer_decorated
|
||||
|
||||
|
||||
class Network(object):
|
||||
|
||||
def __init__(self, inputs, trainable=True):
|
||||
# The input nodes for this network
|
||||
self.inputs = inputs
|
||||
# The current list of terminal nodes
|
||||
self.terminals = []
|
||||
# Mapping from layer names to layers
|
||||
self.layers = dict(inputs)
|
||||
# If true, the resulting variables are set as trainable
|
||||
self.trainable = trainable
|
||||
# Switch variable for dropout
|
||||
self.use_dropout = tf.placeholder_with_default(tf.constant(1.0),
|
||||
shape=[],
|
||||
name='use_dropout')
|
||||
self.setup()
|
||||
|
||||
def setup(self):
|
||||
'''Construct the network. '''
|
||||
raise NotImplementedError('Must be implemented by the subclass.')
|
||||
|
||||
def load(self, data_path, session, ignore_missing=False):
|
||||
'''Load network weights.
|
||||
data_path: The path to the numpy-serialized network weights
|
||||
session: The current TensorFlow session
|
||||
ignore_missing: If true, serialized weights for missing layers are ignored.
|
||||
'''
|
||||
data_dict = np.load(data_path).item()
|
||||
for op_name in data_dict:
|
||||
with tf.variable_scope(op_name, reuse=True):
|
||||
for param_name, data in data_dict[op_name].iteritems():
|
||||
try:
|
||||
var = tf.get_variable(param_name)
|
||||
session.run(var.assign(data))
|
||||
except ValueError:
|
||||
if not ignore_missing:
|
||||
raise
|
||||
|
||||
def feed(self, *args):
|
||||
'''Set the input(s) for the next operation by replacing the terminal nodes.
|
||||
The arguments can be either layer names or the actual layers.
|
||||
'''
|
||||
assert len(args) != 0
|
||||
self.terminals = []
|
||||
for fed_layer in args:
|
||||
if isinstance(fed_layer, basestring):
|
||||
try:
|
||||
fed_layer = self.layers[fed_layer]
|
||||
except KeyError:
|
||||
raise KeyError('Unknown layer name fed: %s' % fed_layer)
|
||||
self.terminals.append(fed_layer)
|
||||
return self
|
||||
|
||||
def get_output(self):
|
||||
'''Returns the current network output.'''
|
||||
return self.terminals[-1]
|
||||
|
||||
def get_unique_name(self, prefix):
|
||||
'''Returns an index-suffixed unique name for the given prefix.
|
||||
This is used for auto-generating layer names based on the type-prefix.
|
||||
'''
|
||||
ident = sum(t.startswith(prefix) for t, _ in self.layers.items()) + 1
|
||||
return '%s_%d' % (prefix, ident)
|
||||
|
||||
def make_var(self, name, shape):
|
||||
'''Creates a new TensorFlow variable.'''
|
||||
return tf.get_variable(name, shape, trainable=self.trainable)
|
||||
|
||||
def validate_padding(self, padding):
|
||||
'''Verifies that the padding is one of the supported ones.'''
|
||||
assert padding in ('SAME', 'VALID')
|
||||
|
||||
@layer
|
||||
def conv(self,
|
||||
input,
|
||||
k_h,
|
||||
k_w,
|
||||
c_o,
|
||||
s_h,
|
||||
s_w,
|
||||
name,
|
||||
relu=True,
|
||||
padding=DEFAULT_PADDING,
|
||||
group=1,
|
||||
biased=True):
|
||||
# Verify that the padding is acceptable
|
||||
self.validate_padding(padding)
|
||||
# Get the number of channels in the input
|
||||
c_i = input.get_shape()[-1]
|
||||
# Verify that the grouping parameter is valid
|
||||
assert c_i % group == 0
|
||||
assert c_o % group == 0
|
||||
# Convolution for a given input and kernel
|
||||
convolve = lambda i, k: tf.nn.conv2d(i, k, [1, s_h, s_w, 1], padding=padding)
|
||||
with tf.variable_scope(name) as scope:
|
||||
kernel = self.make_var('weights', shape=[k_h, k_w, c_i / group, c_o])
|
||||
if group == 1:
|
||||
# This is the common-case. Convolve the input without any further complications.
|
||||
output = convolve(input, kernel)
|
||||
else:
|
||||
# Split the input into groups and then convolve each of them independently
|
||||
input_groups = tf.split(3, group, input)
|
||||
kernel_groups = tf.split(3, group, kernel)
|
||||
output_groups = [convolve(i, k) for i, k in zip(input_groups, kernel_groups)]
|
||||
# Concatenate the groups
|
||||
output = tf.concat(3, output_groups)
|
||||
# Add the biases
|
||||
if biased:
|
||||
biases = self.make_var('biases', [c_o])
|
||||
output = tf.nn.bias_add(output, biases)
|
||||
if relu:
|
||||
# ReLU non-linearity
|
||||
output = tf.nn.relu(output, name=scope.name)
|
||||
return output
|
||||
|
||||
@layer
|
||||
def relu(self, input, name):
|
||||
return tf.nn.relu(input, name=name)
|
||||
|
||||
@layer
|
||||
def max_pool(self, input, k_h, k_w, s_h, s_w, name, padding=DEFAULT_PADDING):
|
||||
self.validate_padding(padding)
|
||||
return tf.nn.max_pool(input,
|
||||
ksize=[1, k_h, k_w, 1],
|
||||
strides=[1, s_h, s_w, 1],
|
||||
padding=padding,
|
||||
name=name)
|
||||
|
||||
@layer
|
||||
def avg_pool(self, input, k_h, k_w, s_h, s_w, name, padding=DEFAULT_PADDING):
|
||||
self.validate_padding(padding)
|
||||
return tf.nn.avg_pool(input,
|
||||
ksize=[1, k_h, k_w, 1],
|
||||
strides=[1, s_h, s_w, 1],
|
||||
padding=padding,
|
||||
name=name)
|
||||
|
||||
@layer
|
||||
def lrn(self, input, radius, alpha, beta, name, bias=1.0):
|
||||
return tf.nn.local_response_normalization(input,
|
||||
depth_radius=radius,
|
||||
alpha=alpha,
|
||||
beta=beta,
|
||||
bias=bias,
|
||||
name=name)
|
||||
|
||||
@layer
|
||||
def concat(self, inputs, axis, name):
|
||||
return tf.concat(concat_dim=axis, values=inputs, name=name)
|
||||
|
||||
@layer
|
||||
def add(self, inputs, name):
|
||||
return tf.add_n(inputs, name=name)
|
||||
|
||||
@layer
|
||||
def fc(self, input, num_out, name, relu=True):
|
||||
with tf.variable_scope(name) as scope:
|
||||
input_shape = input.get_shape()
|
||||
if input_shape.ndims == 4:
|
||||
# The input is spatial. Vectorize it first.
|
||||
dim = 1
|
||||
for d in input_shape[1:].as_list():
|
||||
dim *= d
|
||||
feed_in = tf.reshape(input, [-1, dim])
|
||||
else:
|
||||
feed_in, dim = (input, input_shape[-1].value)
|
||||
weights = self.make_var('weights', shape=[dim, num_out])
|
||||
biases = self.make_var('biases', [num_out])
|
||||
op = tf.nn.relu_layer if relu else tf.nn.xw_plus_b
|
||||
fc = op(feed_in, weights, biases, name=scope.name)
|
||||
return fc
|
||||
|
||||
@layer
|
||||
def softmax(self, input, name):
|
||||
input_shape = map(lambda v: v.value, input.get_shape())
|
||||
if len(input_shape) > 2:
|
||||
# For certain models (like NiN), the singleton spatial dimensions
|
||||
# need to be explicitly squeezed, since they're not broadcast-able
|
||||
# in TensorFlow's NHWC ordering (unlike Caffe's NCHW).
|
||||
if input_shape[1] == 1 and input_shape[2] == 1:
|
||||
input = tf.squeeze(input, squeeze_dims=[1, 2])
|
||||
else:
|
||||
raise ValueError('Rank 2 tensor input expected for softmax!')
|
||||
return tf.nn.softmax(input, name=name)
|
||||
|
||||
@layer
|
||||
def batch_normalization(self, input, name, scale_offset=True, relu=False):
|
||||
# NOTE: Currently, only inference is supported
|
||||
with tf.variable_scope(name) as scope:
|
||||
shape = [input.get_shape()[-1]]
|
||||
if scale_offset:
|
||||
scale = self.make_var('scale', shape=shape)
|
||||
offset = self.make_var('offset', shape=shape)
|
||||
else:
|
||||
scale, offset = (None, None)
|
||||
output = tf.nn.batch_normalization(
|
||||
input,
|
||||
mean=self.make_var('mean', shape=shape),
|
||||
variance=self.make_var('variance', shape=shape),
|
||||
offset=offset,
|
||||
scale=scale,
|
||||
# TODO: This is the default Caffe batch norm eps
|
||||
# Get the actual eps from parameters
|
||||
variance_epsilon=1e-5,
|
||||
name=name)
|
||||
if relu:
|
||||
output = tf.nn.relu(output)
|
||||
return output
|
||||
|
||||
@layer
|
||||
def dropout(self, input, keep_prob, name):
|
||||
keep = 1 - self.use_dropout + (self.use_dropout * keep_prob)
|
||||
return tf.nn.dropout(input, keep, name=name)
|
||||
@@ -0,0 +1,285 @@
|
||||
import numpy as np
|
||||
|
||||
from ..errors import KaffeError, print_stderr
|
||||
from ..graph import GraphBuilder, NodeMapper
|
||||
from ..layers import NodeKind
|
||||
from ..transformers import (DataInjector, DataReshaper, NodeRenamer, ReLUFuser,
|
||||
BatchNormScaleBiasFuser, BatchNormPreprocessor, ParameterNamer)
|
||||
|
||||
from . import network
|
||||
|
||||
|
||||
def get_padding_type(kernel_params, input_shape, output_shape):
|
||||
'''Translates Caffe's numeric padding to one of ('SAME', 'VALID').
|
||||
Caffe supports arbitrary padding values, while TensorFlow only
|
||||
supports 'SAME' and 'VALID' modes. So, not all Caffe paddings
|
||||
can be translated to TensorFlow. There are some subtleties to
|
||||
how the padding edge-cases are handled. These are described here:
|
||||
https://github.com/Yangqing/caffe2/blob/master/caffe2/proto/caffe2_legacy.proto
|
||||
'''
|
||||
k_h, k_w, s_h, s_w, p_h, p_w = kernel_params
|
||||
s_o_h = np.ceil(input_shape.height / float(s_h))
|
||||
s_o_w = np.ceil(input_shape.width / float(s_w))
|
||||
if (output_shape.height == s_o_h) and (output_shape.width == s_o_w):
|
||||
return 'SAME'
|
||||
v_o_h = np.ceil((input_shape.height - k_h + 1.0) / float(s_h))
|
||||
v_o_w = np.ceil((input_shape.width - k_w + 1.0) / float(s_w))
|
||||
if (output_shape.height == v_o_h) and (output_shape.width == v_o_w):
|
||||
return 'VALID'
|
||||
return None
|
||||
|
||||
|
||||
class TensorFlowNode(object):
|
||||
'''An intermediate representation for TensorFlow operations.'''
|
||||
|
||||
def __init__(self, op, *args, **kwargs):
|
||||
# A string corresponding to the TensorFlow operation
|
||||
self.op = op
|
||||
# Positional arguments for the operation
|
||||
self.args = args
|
||||
# Keyword arguments for the operation
|
||||
self.kwargs = list(kwargs.items())
|
||||
# The source Caffe node
|
||||
self.node = None
|
||||
|
||||
def format(self, arg):
|
||||
'''Returns a string representation for the given value.'''
|
||||
return "'%s'" % arg if isinstance(arg, basestring) else str(arg)
|
||||
|
||||
def pair(self, key, value):
|
||||
'''Returns key=formatted(value).'''
|
||||
return '%s=%s' % (key, self.format(value))
|
||||
|
||||
def emit(self):
|
||||
'''Emits the Python source for this node.'''
|
||||
# Format positional arguments
|
||||
args = map(self.format, self.args)
|
||||
# Format any keyword arguments
|
||||
if self.kwargs:
|
||||
args += [self.pair(k, v) for k, v in self.kwargs]
|
||||
# Set the node name
|
||||
args.append(self.pair('name', self.node.name))
|
||||
args = ', '.join(args)
|
||||
return '%s(%s)' % (self.op, args)
|
||||
|
||||
|
||||
class MaybeActivated(object):
|
||||
|
||||
def __init__(self, node, default=True):
|
||||
self.inject_kwargs = {}
|
||||
if node.metadata.get('relu', False) != default:
|
||||
self.inject_kwargs['relu'] = not default
|
||||
|
||||
def __call__(self, *args, **kwargs):
|
||||
kwargs.update(self.inject_kwargs)
|
||||
return TensorFlowNode(*args, **kwargs)
|
||||
|
||||
|
||||
class TensorFlowMapper(NodeMapper):
|
||||
|
||||
def get_kernel_params(self, node):
|
||||
kernel_params = node.layer.kernel_parameters
|
||||
input_shape = node.get_only_parent().output_shape
|
||||
padding = get_padding_type(kernel_params, input_shape, node.output_shape)
|
||||
# Only emit the padding if it's not the default value.
|
||||
padding = {'padding': padding} if padding != network.DEFAULT_PADDING else {}
|
||||
return (kernel_params, padding)
|
||||
|
||||
def map_convolution(self, node):
|
||||
(kernel_params, kwargs) = self.get_kernel_params(node)
|
||||
h = kernel_params.kernel_h
|
||||
w = kernel_params.kernel_w
|
||||
c_o = node.output_shape[1]
|
||||
c_i = node.parents[0].output_shape[1]
|
||||
group = node.parameters.group
|
||||
if group != 1:
|
||||
kwargs['group'] = group
|
||||
if not node.parameters.bias_term:
|
||||
kwargs['biased'] = False
|
||||
assert kernel_params.kernel_h == h
|
||||
assert kernel_params.kernel_w == w
|
||||
return MaybeActivated(node)('conv', kernel_params.kernel_h, kernel_params.kernel_w, c_o,
|
||||
kernel_params.stride_h, kernel_params.stride_w, **kwargs)
|
||||
|
||||
def map_relu(self, node):
|
||||
return TensorFlowNode('relu')
|
||||
|
||||
def map_pooling(self, node):
|
||||
pool_type = node.parameters.pool
|
||||
if pool_type == 0:
|
||||
pool_op = 'max_pool'
|
||||
elif pool_type == 1:
|
||||
pool_op = 'avg_pool'
|
||||
else:
|
||||
# Stochastic pooling, for instance.
|
||||
raise KaffeError('Unsupported pooling type.')
|
||||
(kernel_params, padding) = self.get_kernel_params(node)
|
||||
return TensorFlowNode(pool_op, kernel_params.kernel_h, kernel_params.kernel_w,
|
||||
kernel_params.stride_h, kernel_params.stride_w, **padding)
|
||||
|
||||
def map_inner_product(self, node):
|
||||
#TODO: Axis
|
||||
assert node.parameters.axis == 1
|
||||
#TODO: Unbiased
|
||||
assert node.parameters.bias_term == True
|
||||
return MaybeActivated(node)('fc', node.parameters.num_output)
|
||||
|
||||
def map_softmax(self, node):
|
||||
return TensorFlowNode('softmax')
|
||||
|
||||
def map_lrn(self, node):
|
||||
params = node.parameters
|
||||
# The window size must be an odd value. For a window
|
||||
# size of (2*n+1), TensorFlow defines depth_radius = n.
|
||||
assert params.local_size % 2 == 1
|
||||
# Caffe scales by (alpha/(2*n+1)), whereas TensorFlow
|
||||
# just scales by alpha (as does Krizhevsky's paper).
|
||||
# We'll account for that here.
|
||||
alpha = params.alpha / float(params.local_size)
|
||||
return TensorFlowNode('lrn', int(params.local_size / 2), alpha, params.beta)
|
||||
|
||||
def map_concat(self, node):
|
||||
axis = (2, 3, 1, 0)[node.parameters.axis]
|
||||
return TensorFlowNode('concat', axis)
|
||||
|
||||
def map_dropout(self, node):
|
||||
return TensorFlowNode('dropout', node.parameters.dropout_ratio)
|
||||
|
||||
def map_batch_norm(self, node):
|
||||
scale_offset = len(node.data) == 4
|
||||
kwargs = {} if scale_offset else {'scale_offset': False}
|
||||
return MaybeActivated(node, default=False)('batch_normalization', **kwargs)
|
||||
|
||||
def map_eltwise(self, node):
|
||||
operations = {0: 'multiply', 1: 'add', 2: 'max'}
|
||||
op_code = node.parameters.operation
|
||||
try:
|
||||
return TensorFlowNode(operations[op_code])
|
||||
except KeyError:
|
||||
raise KaffeError('Unknown elementwise operation: {}'.format(op_code))
|
||||
|
||||
def commit(self, chains):
|
||||
return chains
|
||||
|
||||
|
||||
class TensorFlowEmitter(object):
|
||||
|
||||
def __init__(self, tab=None):
|
||||
self.tab = tab or ' ' * 4
|
||||
self.prefix = ''
|
||||
|
||||
def indent(self):
|
||||
self.prefix += self.tab
|
||||
|
||||
def outdent(self):
|
||||
self.prefix = self.prefix[:-len(self.tab)]
|
||||
|
||||
def statement(self, s):
|
||||
return self.prefix + s + '\n'
|
||||
|
||||
def emit_imports(self):
|
||||
return self.statement('from kaffe.tensorflow import Network\n')
|
||||
|
||||
def emit_class_def(self, name):
|
||||
return self.statement('class %s(Network):' % (name))
|
||||
|
||||
def emit_setup_def(self):
|
||||
return self.statement('def setup(self):')
|
||||
|
||||
def emit_parents(self, chain):
|
||||
assert len(chain)
|
||||
s = '(self.feed('
|
||||
sep = ', \n' + self.prefix + (' ' * len(s))
|
||||
s += sep.join(["'%s'" % parent.name for parent in chain[0].node.parents])
|
||||
return self.statement(s + ')')
|
||||
|
||||
def emit_node(self, node):
|
||||
return self.statement(' ' * 5 + '.' + node.emit())
|
||||
|
||||
def emit(self, name, chains):
|
||||
s = self.emit_imports()
|
||||
s += self.emit_class_def(name)
|
||||
self.indent()
|
||||
s += self.emit_setup_def()
|
||||
self.indent()
|
||||
blocks = []
|
||||
for chain in chains:
|
||||
b = ''
|
||||
b += self.emit_parents(chain)
|
||||
for node in chain:
|
||||
b += self.emit_node(node)
|
||||
blocks.append(b[:-1] + ')')
|
||||
s = s + '\n\n'.join(blocks)
|
||||
return s
|
||||
|
||||
|
||||
class TensorFlowTransformer(object):
|
||||
|
||||
def __init__(self, def_path, data_path, verbose=True, phase='test'):
|
||||
self.verbose = verbose
|
||||
self.phase = phase
|
||||
self.load(def_path, data_path, phase)
|
||||
self.params = None
|
||||
self.source = None
|
||||
|
||||
def load(self, def_path, data_path, phase):
|
||||
# Build the graph
|
||||
graph = GraphBuilder(def_path, phase).build()
|
||||
|
||||
if data_path is not None:
|
||||
# Load and associate learned parameters
|
||||
graph = DataInjector(def_path, data_path)(graph)
|
||||
|
||||
# Transform the graph
|
||||
transformers = [
|
||||
# Fuse split batch normalization layers
|
||||
BatchNormScaleBiasFuser(),
|
||||
|
||||
# Fuse ReLUs
|
||||
# TODO: Move non-linearity application to layer wrapper, allowing
|
||||
# any arbitrary operation to be optionally activated.
|
||||
ReLUFuser(allowed_parent_types=[NodeKind.Convolution, NodeKind.InnerProduct,
|
||||
NodeKind.BatchNorm]),
|
||||
|
||||
# Rename nodes
|
||||
# Slashes are used for scoping in TensorFlow. Replace slashes
|
||||
# in node names with underscores.
|
||||
# (Caffe's GoogLeNet implementation uses slashes)
|
||||
NodeRenamer(lambda node: node.name.replace('/', '_'))
|
||||
]
|
||||
self.graph = graph.transformed(transformers)
|
||||
|
||||
# Display the graph
|
||||
if self.verbose:
|
||||
print_stderr(self.graph)
|
||||
|
||||
def transform_data(self):
|
||||
if self.params is None:
|
||||
transformers = [
|
||||
|
||||
# Reshape the parameters to TensorFlow's ordering
|
||||
DataReshaper({
|
||||
# (c_o, c_i, h, w) -> (h, w, c_i, c_o)
|
||||
NodeKind.Convolution: (2, 3, 1, 0),
|
||||
|
||||
# (c_o, c_i) -> (c_i, c_o)
|
||||
NodeKind.InnerProduct: (1, 0)
|
||||
}),
|
||||
|
||||
# Pre-process batch normalization data
|
||||
BatchNormPreprocessor(),
|
||||
|
||||
# Convert parameters to dictionaries
|
||||
ParameterNamer(),
|
||||
]
|
||||
self.graph = self.graph.transformed(transformers)
|
||||
self.params = {node.name: node.data for node in self.graph.nodes if node.data}
|
||||
return self.params
|
||||
|
||||
def transform_source(self):
|
||||
if self.source is None:
|
||||
mapper = TensorFlowMapper(self.graph)
|
||||
chains = mapper.map()
|
||||
emitter = TensorFlowEmitter()
|
||||
self.source = emitter.emit(self.graph.name, chains)
|
||||
return self.source
|
||||
Reference in New Issue
Block a user