mirror of
https://github.com/wassname/PSPNet-Keras-tensorflow.git
synced 2026-09-10 11:40:16 +08:00
Made commit of original caffe-tensorflow converter
This commit is contained in:
VladK
@@ -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)
|
||||
Reference in New Issue
Block a user