def Layernorm(name, norm_axes, inputs):
mean, var = tf.nn.moments(inputs, norm_axes, keep_dims=True)
# Assume the 'neurons' axis is the first of norm_axes. This is the case for fully-connected and BCHW conv layers.
n_neurons = inputs.get_shape().as_list()[norm_axes[0]]
offset = lib.param(name + '.offset', np.zeros(n_neurons, dtype='float32'))
scale = lib.param(name + '.scale', np.ones(n_neurons, dtype='float32'))
# Add broadcasting dims to offset and scale (e.g. BCHW conv data)
offset = tf.reshape(offset, [-1] + [1 for i in xrange(len(norm_axes) - 1)])
scale = tf.reshape(scale, [-1] + [1 for i in xrange(len(norm_axes) - 1)])
result = tf.nn.batch_normalization(inputs, mean, var, offset, scale, 1e-5)
return result
def Batchnorm(name,
axes,
inputs,
is_training=None,
stats_iter=None,
update_moving_stats=True,
fused=True,
labels=None,
n_labels=None):
"""conditional batchnorm (dumoulin et al 2016) for BCHW conv filtermaps"""
if axes != [0, 2, 3]:
raise Exception('unsupported')
mean, var = tf.nn.moments(inputs, axes, keep_dims=True)
shape = mean.get_shape().as_list() # shape is [1,n,1,1]
offset_m = lib.param(name + '.offset',
np.zeros([n_labels, shape[1]], dtype='float32'))
scale_m = lib.param(name + '.scale',
np.ones([n_labels, shape[1]], dtype='float32'))
offset = tf.nn.embedding_lookup(offset_m, labels)
scale = tf.nn.embedding_lookup(scale_m, labels)
result = tf.nn.batch_normalization(inputs, mean, var, offset[:, :, None,
None],
scale[:, :, None, None], 1e-5)
return result
def Linear(
name,
input_dim,
output_dim,
inputs,
biases=True,
initialization=None,
weightnorm=None,
gain=1.
):
"""
initialization: None, `lecun`, 'glorot', `he`, 'glorot_he', `orthogonal`, `("uniform", range)`
"""
with tf.name_scope(name) as scope:
def uniform(stdev, size):
if _weights_stdev is not None:
stdev = _weights_stdev
return np.random.uniform(
low=-stdev * np.sqrt(3),
high=stdev * np.sqrt(3),
size=size
).astype('float32')
if initialization == 'lecun':# and input_dim != output_dim):
# disabling orth. init for now because it's too slow
weight_values = uniform(
np.sqrt(1./input_dim),
(input_dim, output_dim)
)
elif initialization == 'glorot' or (initialization == None):
weight_values = uniform(
np.sqrt(2./(input_dim+output_dim)),
(input_dim, output_dim)
)
elif initialization == 'he':
weight_values = uniform(
np.sqrt(2./input_dim),
(input_dim, output_dim)
)
elif initialization == 'glorot_he':
weight_values = uniform(
np.sqrt(4./(input_dim+output_dim)),
(input_dim, output_dim)
)
elif initialization == 'orthogonal' or \
(initialization == None and input_dim == output_dim):
# From lasagne
def sample(shape):
if len(shape) < 2:
raise RuntimeError("Only shapes of length 2 or more are "
"supported.")
flat_shape = (shape[0], np.prod(shape[1:]))
# TODO: why normal and not uniform?
a = np.random.normal(0.0, 1.0, flat_shape)
u, _, v = np.linalg.svd(a, full_matrices=False)
# pick the one with the correct shape
q = u if u.shape == flat_shape else v
q = q.reshape(shape)
return q.astype('float32')
weight_values = sample((input_dim, output_dim))
elif initialization[0] == 'uniform':
weight_values = np.random.uniform(
low=-initialization[1],
high=initialization[1],
size=(input_dim, output_dim)
).astype('float32')
else:
raise Exception('Invalid initialization!')
weight_values *= gain
weight = lib.param(
name + '.W',
weight_values
)
if weightnorm==None:
weightnorm = _default_weightnorm
if weightnorm:
norm_values = np.sqrt(np.sum(np.square(weight_values), axis=0))
# norm_values = np.linalg.norm(weight_values, axis=0)
target_norms = lib.param(
name + '.g',
norm_values
)
with tf.name_scope('weightnorm') as scope:
norms = tf.sqrt(tf.reduce_sum(tf.square(weight), reduction_indices=[0]))
weight = weight * (target_norms / norms)
# if 'Discriminator' in name:
# print "WARNING weight constraint on {}".format(name)
# weight = tf.nn.softsign(10.*weight)*.1
if inputs.get_shape().ndims == 2:
result = tf.matmul(inputs, weight)
else:
reshaped_inputs = tf.reshape(inputs, [-1, input_dim])
result = tf.matmul(reshaped_inputs, weight)
result = tf.reshape(result, tf.pack(tf.unpack(tf.shape(inputs))[:-1] + [output_dim]))
if biases:
result = tf.nn.bias_add(
result,
lib.param(
name + '.b',
np.zeros((output_dim,), dtype='float32')
)
)
return result
def Conv1D(name, input_dim, output_dim, filter_size, inputs, he_init=True, mask_type=None, stride=1, weightnorm=None, biases=True, gain=1.):
"""
inputs: tensor of shape (batch size, num channels, width)
mask_type: one of None, 'a', 'b'
returns: tensor of shape (batch size, num channels, width)
"""
with tf.name_scope(name) as scope:
if mask_type is not None:
mask_type, mask_n_channels = mask_type
mask = np.ones(
(filter_size, input_dim, output_dim),
dtype='float32'
)
center = filter_size // 2
# Mask out future locations
# filter shape is (width, input channels, output channels)
mask[center+1:, :, :] = 0.
# Mask out future channels
for i in xrange(mask_n_channels):
for j in xrange(mask_n_channels):
if (mask_type=='a' and i >= j) or (mask_type=='b' and i > j):
mask[
center,
i::mask_n_channels,
j::mask_n_channels
] = 0.
def uniform(stdev, size):
return np.random.uniform(
low=-stdev * np.sqrt(3),
high=stdev * np.sqrt(3),
size=size
).astype('float32')
fan_in = input_dim * filter_size
fan_out = output_dim * filter_size / stride
if mask_type is not None: # only approximately correct
fan_in /= 2.
fan_out /= 2.
if he_init:
filters_stdev = np.sqrt(4./(fan_in+fan_out))
else: # Normalized init (Glorot & Bengio)
filters_stdev = np.sqrt(2./(fan_in+fan_out))
filter_values = uniform(
filters_stdev,
(filter_size, input_dim, output_dim)
)
# print "WARNING IGNORING GAIN"
filter_values *= gain
filters = lib.param(name+'.Filters', filter_values)
if weightnorm==None:
weightnorm = _default_weightnorm
if weightnorm:
norm_values = np.sqrt(np.sum(np.square(filter_values), axis=(0,1)))
target_norms = lib.param(
name + '.g',
norm_values
)
with tf.name_scope('weightnorm') as scope:
norms = tf.sqrt(tf.reduce_sum(tf.square(filters), reduction_indices=[0,1]))
filters = filters * (target_norms / norms)
if mask_type is not None:
with tf.name_scope('filter_mask'):
filters = filters * mask
result = tf.nn.conv1d(
value=inputs,
filters=filters,
stride=stride,
padding='SAME',
data_format='NCHW'
)
if biases:
_biases = lib.param(
name+'.Biases',
np.zeros([output_dim], dtype='float32')
)
# result = result + _biases
result = tf.expand_dims(result, 3)
result = tf.nn.bias_add(result, _biases, data_format='NCHW')
result = tf.squeeze(result)
return result
def Deconv2D(
name,
input_dim,
output_dim,
filter_size,
inputs,
he_init=True,
weightnorm=None,
biases=True,
gain=1.,
mask_type=None,
):
"""
inputs: tensor of shape (batch size, height, width, input_dim)
returns: tensor of shape (batch size, 2*height, 2*width, output_dim)
"""
with tf.name_scope(name) as scope:
if mask_type != None:
raise Exception('Unsupported configuration')
def uniform(stdev, size):
return np.random.uniform(
low=-stdev * np.sqrt(3),
high=stdev * np.sqrt(3),
size=size
).astype('float32')
stride = 2
fan_in = input_dim * filter_size**2 / (stride**2)
fan_out = output_dim * filter_size**2
if he_init:
filters_stdev = np.sqrt(4./(fan_in+fan_out))
else: # Normalized init (Glorot & Bengio)
filters_stdev = np.sqrt(2./(fan_in+fan_out))
if _weights_stdev is not None:
filter_values = uniform(
_weights_stdev,
(filter_size, filter_size, output_dim, input_dim)
)
else:
filter_values = uniform(
filters_stdev,
(filter_size, filter_size, output_dim, input_dim)
)
filter_values *= gain
filters = lib.param(
name+'.Filters',
filter_values
)
if weightnorm==None:
weightnorm = _default_weightnorm
if weightnorm:
norm_values = np.sqrt(np.sum(np.square(filter_values), axis=(0,1,3)))
target_norms = lib.param(
name + '.g',
norm_values
)
with tf.name_scope('weightnorm') as scope:
norms = tf.sqrt(tf.reduce_sum(tf.square(filters), reduction_indices=[0,1,3]))
filters = filters * tf.expand_dims(target_norms / norms, 1)
inputs = tf.transpose(inputs, [0,2,3,1], name='NCHW_to_NHWC')
input_shape = tf.shape(inputs)
try: # tf pre-1.0 (top) vs 1.0 (bottom)
output_shape = tf.pack([input_shape[0], 2*input_shape[1], 2*input_shape[2], output_dim])
except Exception as e:
output_shape = tf.stack([input_shape[0], 2*input_shape[1], 2*input_shape[2], output_dim])
result = tf.nn.conv2d_transpose(
value=inputs,
filter=filters,
output_shape=output_shape,
strides=[1, 2, 2, 1],
padding='SAME'
)
if biases:
_biases = lib.param(
name+'.Biases',
np.zeros(output_dim, dtype='float32')
)
result = tf.nn.bias_add(result, _biases)
result = tf.transpose(result, [0,3,1,2], name='NHWC_to_NCHW')
return result
def Batchnorm(name,
axes,
inputs,
is_training=None,
stats_iter=None,
update_moving_stats=True,
fused=True):
if ((axes == [0, 2, 3]) or (axes == [0, 2])) and fused == True:
if axes == [0, 2]:
inputs = tf.expand_dims(inputs, 3)
# Old (working but pretty slow) implementation:
##########
# inputs = tf.transpose(inputs, [0,2,3,1])
# mean, var = tf.nn.moments(inputs, [0,1,2], keep_dims=False)
# offset = lib.param(name+'.offset', np.zeros(mean.get_shape()[-1], dtype='float32'))
# scale = lib.param(name+'.scale', np.ones(var.get_shape()[-1], dtype='float32'))
# result = tf.nn.batch_normalization(inputs, mean, var, offset, scale, 1e-4)
# return tf.transpose(result, [0,3,1,2])
# New (super fast but untested) implementation:
offset = lib.param(name + '.offset',
np.zeros(inputs.get_shape()[1], dtype='float32'))
scale = lib.param(name + '.scale',
np.ones(inputs.get_shape()[1], dtype='float32'))
moving_mean = lib.param(name + '.moving_mean',
np.zeros(inputs.get_shape()[1],
dtype='float32'),
trainable=False)
moving_variance = lib.param(name + '.moving_variance',
np.ones(inputs.get_shape()[1],
dtype='float32'),
trainable=False)
def _fused_batch_norm_training():
return tf.nn.fused_batch_norm(inputs,
scale,
offset,
epsilon=1e-5,
data_format='NCHW')
def _fused_batch_norm_inference():
# Version which blends in the current item's statistics
batch_size = tf.cast(tf.shape(inputs)[0], 'float32')
mean, var = tf.nn.moments(inputs, [2, 3], keep_dims=True)
mean = ((1. / batch_size) * mean) + ((
(batch_size - 1.) / batch_size) * moving_mean)[None, :, None,
None]
var = ((1. / batch_size) * var) + ((
(batch_size - 1.) / batch_size) * moving_variance)[None, :,
None, None]
return tf.nn.batch_normalization(inputs, mean, var,
offset[None, :, None, None],
scale[None, :, None,
None], 1e-5), mean, var
# Standard version
# return tf.nn.fused_batch_norm(
# inputs,
# scale,
# offset,
# epsilon=1e-2,
# mean=moving_mean,
# variance=moving_variance,
# is_training=False,
# data_format='NCHW'
# )
if is_training is None:
outputs, batch_mean, batch_var = _fused_batch_norm_training()
else:
outputs, batch_mean, batch_var = tf.cond(
is_training, _fused_batch_norm_training,
_fused_batch_norm_inference)
if update_moving_stats:
no_updates = lambda: outputs
def _force_updates():
"""Internal function forces updates moving_vars if is_training."""
float_stats_iter = tf.cast(stats_iter, tf.float32)
update_moving_mean = tf.assign(
moving_mean, ((float_stats_iter /
(float_stats_iter + 1)) * moving_mean) +
((1 / (float_stats_iter + 1)) * batch_mean))
update_moving_variance = tf.assign(
moving_variance,
((float_stats_iter /
(float_stats_iter + 1)) * moving_variance) +
((1 / (float_stats_iter + 1)) * batch_var))
with tf.control_dependencies(
[update_moving_mean, update_moving_variance]):
return tf.identity(outputs)
outputs = tf.cond(is_training, _force_updates, no_updates)
if axes == [0, 2]:
return outputs[:, :, :, 0] # collapse last dim
else:
return outputs
else:
# raise Exception('old BN')
# TODO we can probably use nn.fused_batch_norm here too for speedup
mean, var = tf.nn.moments(inputs, axes, keep_dims=True)
shape = mean.get_shape().as_list()
if 0 not in axes:
print(
"WARNING ({}): didn't find 0 in axes, but not using separate BN params for each item in batch"
.format(name))
shape[0] = 1
offset = lib.param(name + '.offset', np.zeros(shape, dtype='float32'))
scale = lib.param(name + '.scale', np.ones(shape, dtype='float32'))
result = tf.nn.batch_normalization(inputs, mean, var, offset, scale,
1e-5)
return result