def mpusim_fully_connected(inputs,
units,
activation=None,
use_bias=True,
kernel_initializer=None,
bias_initializer=tf.zeros_initializer(),
kernel_regularizer=None,
bias_regularizer=None,
activity_regularizer=None,
activations_datatype_size_byte=1,
weights_datatype_size_byte=1,
results_datatype_size_byte=4,
systolic_array_height=256,
systolic_array_width=256,
activation_fifo_depth=8,
accumulator_array_height=4096,
log_file_output_dir='.',
model_name='unnamed'):
"""
A wrapper around `mpusim_fc`.
One difference to maintain backward-compatibility:
Default weight initializer is variance_scaling_initializer(2.0).
Variable Names:
* ``W``: weights of shape [in_dim, out_dim]
* ``b``: bias
"""
if kernel_initializer is None:
if get_tf_version_tuple() <= (1, 12):
kernel_initializer = tf.contrib.layers.variance_scaling_initializer(2.0) # deprecated
else:
kernel_initializer = tf.keras.initializers.VarianceScaling(2.0, distribution='untruncated_normal')
inputs = batch_flatten(inputs)
with rename_get_variable({'kernel': 'W', 'bias': 'b'}):
layer = mpusim_fc(units=units,
activation=activation,
use_bias=use_bias,
kernel_initializer=kernel_initializer,
bias_initializer=bias_initializer,
kernel_regularizer=kernel_regularizer,
bias_regularizer=bias_regularizer,
activity_regularizer=activity_regularizer,
activations_datatype_size_byte=activations_datatype_size_byte,
weights_datatype_size_byte=weights_datatype_size_byte,
results_datatype_size_byte=results_datatype_size_byte,
systolic_array_height=systolic_array_height,
systolic_array_width=systolic_array_width,
activation_fifo_depth=activation_fifo_depth,
accumulator_array_height=accumulator_array_height,
log_file_output_dir=log_file_output_dir,
model_name=model_name,
_reuse=tf.get_variable_scope().reuse)
ret = layer.apply(inputs, scope=tf.get_variable_scope())
ret = tf.identity(ret, name='output')
ret.variables = VariableHolder(W=layer.kernel)
if use_bias:
ret.variables.b = layer.bias
return ret
def mod_conv2d(x,
y,
fmaps,
kernel,
demodulate=True,
gain=1,
use_wscale=True,
lrmul=1,
fused_modconv=True,
eps=1e-8,
padding='SAME',
name="mod_conv2d"):
shape = x.get_shape().as_list() # [n, h, w, c]
cin = shape[-1]
with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
# Get weight
w = get_weight([kernel, kernel, cin, fmaps],
gain=gain,
use_wscale=use_wscale,
lrmul=lrmul,
name='W')
ww = w[tf.newaxis] # introduce minibatch dimension
# Modulate
s = get_bias(
cin, base_std=0, use_wscale=use_wscale, lrmul=lrmul, name='bs') + 1
vh = VariableHolder(W=w, bs=s)
s = tf.tile(s[tf.newaxis],
[tf.shape(x)[0], 1]) # introduce minibatch dimension
if y is not None:
y_style, w_style = dense(y,
cin,
gain=gain,
use_wscale=use_wscale,
lrmul=lrmul)
s = s + y_style
vh.Ws = w_style
ww = ww * tf.cast(s[:, tf.newaxis, tf.newaxis, :, tf.newaxis],
w.dtype) # scale input feature maps
# Demodulate
if demodulate:
d = tf.rsqrt(
tf.reduce_sum(tf.square(ww), axis=[1, 2, 3], keepdims=True) +
eps) # scaling factor
ww = ww * d
# Reshape/scale input
if fused_modconv:
x = tf.reshape(tf.transpose(x, [0, 3, 1, 2]),
[1, -1, shape[1], shape[2]]) # [1, n*cin, h, w]
w = tf.reshape(tf.transpose(ww, [1, 2, 3, 0, 4]),
[kernel, kernel, cin, -1]) # [k, k, cin, n*cout]
x = tf.nn.conv2d(x,
tf.cast(w, x.dtype),
data_format='NCHW',
strides=[1, 1, 1, 1],
padding=padding)
out_shape = x.get_shape().as_list()
x = tf.transpose(
tf.reshape(x, [-1, fmaps, out_shape[2], out_shape[3]]),
[0, 2, 3, 1])
else:
x = x * tf.cast(s[:, tf.newaxis, tf.newaxis, :], x.dtype)
x = tf.nn.conv2d(x,
tf.cast(w, x.dtype),
data_format='NHWC',
strides=[1, 1, 1, 1],
padding=padding)
if demodulate:
x = x * tf.cast(tf.reshape(d, [-1, 1, 1, fmaps]), x.dtype)
ret = tf.identity(x)
ret.variables = vh
return ret
def Conv(inputs,
filters,
kernel_size,
strides=(1, 1),
padding='same',
data_format='channels_last',
dilation_rate=(1, 1),
activation=None,
use_bias=True,
kernel_initializer=None,
bias_initializer=tf.zeros_initializer(),
kernel_regularizer=None,
bias_regularizer=None,
activity_regularizer=None,
split=1,
norm=False):
"""
Similar to `tf.layers.Conv2D`, but with some differences:
1. Default kernel initializer is variance_scaling_initializer(2.0).
2. Default padding is 'same'.
3. Support 'split' argument to do group convolution.
Variable Names:
* ``W``: weights
* ``b``: bias
"""
if kernel_initializer is None:
if get_tf_version_tuple() <= (1, 12):
kernel_initializer = tf.contrib.layers.variance_scaling_initializer(
2.0) # deprecated
else:
kernel_initializer = tf.keras.initializers.VarianceScaling(
2.0, distribution='untruncated_normal')
dilation_rate = shape2d(dilation_rate)
if True:
# group conv implementation
data_format = get_data_format(data_format, keras_mode=False)
in_shape = inputs.get_shape().as_list()
channel_axis = 3 if data_format == 'NHWC' else 1
in_channel = in_shape[channel_axis]
assert in_channel is not None, "[Conv2D] Input cannot have unknown channel!"
assert in_channel % split == 0
assert kernel_regularizer is None and bias_regularizer is None and activity_regularizer is None, \
"Not supported by group conv or dilated conv!"
out_channel = filters
assert out_channel % split == 0
assert dilation_rate == [1, 1] or get_tf_version_tuple() >= (
1, 5), 'TF>=1.5 required for dilated conv.'
kernel_shape = shape2d(kernel_size)
filter_shape = kernel_shape + [in_channel // split, out_channel]
stride = shape4d(strides, data_format=data_format)
kwargs = {"data_format": data_format}
if get_tf_version_tuple() >= (1, 5):
kwargs['dilations'] = shape4d(dilation_rate,
data_format=data_format)
# matching input dtype (ex. tf.float16) since the default dtype of variable if tf.float32
inputs_dtype = inputs.dtype
W = tf.get_variable('parseweigth',
filter_shape,
dtype=inputs_dtype,
initializer=kernel_initializer)
if norm:
use_bias = False
W = tf.reshape(W, kernel_shape + [4, in_channel // 4, out_channel])
W = tf.nn.softmax(W, 2)
W = tf.reshape(W, filter_shape)
#dynamics = tf.reduce_mean(inputs, 0)
#dynamics = tf.transpose(dynamics, [1,2,0])
#dynamics = tf.image.resize_images(dynamics, kernel_shape)
#dynamics = tf.expand_dims(dynamics, -1)
#W = W + 0.001 * dynamics #tf.random_normal(shape = tf.shape(W), mean = 0.0, stddev = 0.012, dtype = tf.float32)
#W = W *tf.random_uniform(shape=W.get_shape().as_list(), minval=0., maxval=2.)
if use_bias:
b = tf.get_variable('parsebias', [out_channel],
dtype=inputs_dtype,
initializer=bias_initializer)
if split == 1:
conv = tf.nn.conv2d(inputs, W, stride, padding.upper(), **kwargs)
else:
try:
conv = tf.nn.conv2d(inputs, W, stride, padding.upper(),
**kwargs)
except ValueError:
log_once(
"CUDNN group convolution support is only available with "
"https://github.com/tensorflow/tensorflow/pull/25818 . "
"Will fall back to a loop-based slow implementation instead!",
'warn')
ret = tf.nn.bias_add(conv, b,
data_format=data_format) if use_bias else conv
if activation is not None:
ret = activation(ret)
ret = tf.identity(ret, name='output')
ret.variables = VariableHolder(W=W)
if use_bias:
ret.variables.b = b
return ret
def mpusim_depthwise_convolution2d(inputs,
kernel_size,
strides=(1, 1),
padding='valid',
depth_multiplier=1,
data_format='channels_last',
activation=None,
use_bias=True,
depthwise_initializer='glorot_uniform',
bias_initializer='zeros',
depthwise_regularizer=None,
bias_regularizer=None,
depthwise_constraint=None,
bias_constraint=None,
activations_datatype_size_byte=1,
weights_datatype_size_byte=1,
results_datatype_size_byte=4,
systolic_array_height=256,
systolic_array_width=256,
activation_fifo_depth=8,
accumulator_array_height=4096,
log_file_output_dir='.',
model_name='unnamed'):
#depthwise_initializer = initializers.get(depthwise_initializer)
#depthwise_regularizer = regularizers.get(depthwise_regularizer)
#depthwise_constraint = constraints.get(depthwise_constraint)
#bias_initializer = initializers.get(bias_initializer)
data_format = get_data_format(data_format, keras_mode=False)
input_shape = inputs.get_shape().as_list()
strides = shape4d(strides, data_format=data_format)
if len(input_shape) < 4:
raise ValueError(
'Inputs to `mpusim_depthwise_conv2d` should have rank 4. '
'Received input shape:', str(input_shape))
if data_format == 'NCHW':
raise ValueError('mpusim_depthwise_convolution2d '
'only supports NHWC data format')
else:
channel_axis = 3
if input_shape[channel_axis] is None:
raise ValueError('The channel dimension of the inputs to '
'`mpusim_depthwise_convolution2d` '
'should be defined. Found `None`.')
input_dim = int(input_shape[channel_axis])
depthwise_kernel_shape = (kernel_size[0], kernel_size[1], input_dim,
depth_multiplier)
depthwise_kernel = tf.get_variable('W',
shape=depthwise_kernel_shape,
initializer=depthwise_initializer,
regularizer=depthwise_regularizer,
constraint=depthwise_constraint)
if use_bias:
biases = tf.get_variable('b',
shape=(input_dim * depth_multiplier, ),
initializer=bias_initializer,
regularizer=bias_regularizer,
constraint=bias_constraint)
result = mpusim_depthwise_conv2d(
inputs,
depthwise_kernel,
strides=strides,
padding=padding,
data_format=data_format,
activations_datatype_size_byte=activations_datatype_size_byte,
weights_datatype_size_byte=weights_datatype_size_byte,
results_datatype_size_byte=results_datatype_size_byte,
systolic_array_height=systolic_array_height,
systolic_array_width=systolic_array_width,
activation_fifo_depth=activation_fifo_depth,
accumulator_array_height=accumulator_array_height,
log_file_output_dir=log_file_output_dir,
model_name=model_name)
if use_bias:
result = tf.nn.bias_add(result, bias, data_format=data_format)
if activation is not None:
result = activation(result)
result = tf.identity(result, name='output')
result.variables = VariableHolder(W=depthwise_kernel)
if use_bias:
result.variables.b = biases
return result
def MaskedConv2D(
inputs,
filters,
kernel_size,
strides=(1, 1),
padding='same',
data_format='channels_last',
dilation_rate=(1, 1),
activation=None,
use_bias=True,
kernel_initializer=None,
bias_initializer=tf.zeros_initializer(),
kernel_regularizer=None,
bias_regularizer=None,
activity_regularizer=None,
split=1,
masking=False):
"""
A wrapper around `tf.layers.Conv2D`.
Some differences to maintain backward-compatibility:
1. Default kernel initializer is variance_scaling_initializer(2.0).
2. Default padding is 'same'.
3. Support 'split' argument to do group conv.
Variable Names:
* ``W``: weights
* ``b``: bias
"""
if kernel_initializer is None:
if get_tf_version_tuple() <= (1, 12):
kernel_initializer = tf.contrib.layers.variance_scaling_initializer(2.0)
else:
kernel_initializer = tf.keras.initializers.VarianceScaling(2.0, distribution='untruncated_normal')
dilation_rate = shape2d(dilation_rate)
if (masking == False) and (split == 1) and (dilation_rate == [1, 1]):
# tf.layers.Conv2D has bugs with dilations (https://github.com/tensorflow/tensorflow/issues/26797)
with rename_get_variable({'kernel': 'W', 'bias': 'b'}):
layer = tf.layers.Conv2D(
filters,
kernel_size,
strides=strides,
padding=padding,
data_format=data_format,
dilation_rate=dilation_rate,
activation=activation,
use_bias=use_bias,
kernel_initializer=kernel_initializer,
bias_initializer=bias_initializer,
kernel_regularizer=kernel_regularizer,
bias_regularizer=bias_regularizer,
activity_regularizer=activity_regularizer,
_reuse=tf.get_variable_scope().reuse)
ret = layer.apply(inputs, scope=tf.get_variable_scope())
ret = tf.identity(ret, name='output')
ret.variables = VariableHolder(W=layer.kernel)
if use_bias:
ret.variables.b = layer.bias
else:
if masking == True:
assert split == 1, "Pruining group conv is not supported yet"
# group conv implementation
data_format = get_data_format(data_format, keras_mode=False)
in_shape = inputs.get_shape().as_list()
channel_axis = 3 if data_format == 'NHWC' else 1
in_channel = in_shape[channel_axis]
assert in_channel is not None, "[Conv2D] Input cannot have unknown channel!"
assert in_channel % split == 0
assert kernel_regularizer is None and bias_regularizer is None and activity_regularizer is None, \
"Not supported by group conv or dilated conv!"
out_channel = filters
assert out_channel % split == 0
assert dilation_rate == [1, 1] or get_tf_version_tuple() >= (1, 5), 'TF>=1.5 required for dilated conv.'
kernel_shape = shape2d(kernel_size)
filter_shape = kernel_shape + [in_channel / split, out_channel]
stride = shape4d(strides, data_format=data_format)
kwargs = dict(data_format=data_format)
if get_tf_version_tuple() >= (1, 5):
kwargs['dilations'] = shape4d(dilation_rate, data_format=data_format)
W = tf.get_variable(
'W', filter_shape, initializer=kernel_initializer)
if use_bias:
b = tf.get_variable('b', [out_channel], initializer=bias_initializer)
if split == 1:
if masking:
W = pruning.apply_mask(W)
conv = tf.nn.conv2d(inputs, W, stride, padding.upper(), **kwargs)
else:
conv = None
if get_tf_version_tuple() >= (1, 13):
try:
conv = tf.nn.conv2d(inputs, W, stride, padding.upper(), **kwargs)
except ValueError:
log_once("CUDNN group convolution support is only available with "
"https://github.com/tensorflow/tensorflow/pull/25818 . "
"Will fall back to a loop-based slow implementation instead!", 'warn')
if conv is None:
inputs = tf.split(inputs, split, channel_axis)
kernels = tf.split(W, split, 3)
outputs = [tf.nn.conv2d(i, k, stride, padding.upper(), **kwargs)
for i, k in zip(inputs, kernels)]
conv = tf.concat(outputs, channel_axis)
ret = tf.nn.bias_add(conv, b, data_format=data_format) if use_bias else conv
if activation is not None:
ret = activation(ret)
ret = tf.identity(ret, name='output')
ret.variables = VariableHolder(W=W)
if use_bias:
ret.variables.b = b
return ret
def mpusim_conv2d(
inputs,
filters,
kernel_size,
strides=(1, 1),
padding='same',
data_format='channels_last',
dilation_rate=(1, 1),
activation=None,
use_bias=True,
kernel_initializer=None,
bias_initializer=tf.zeros_initializer(),
kernel_regularizer=None,
bias_regularizer=None,
activity_regularizer=None,
split=1,
activations_datatype_size_byte=1,
weights_datatype_size_byte=1,
results_datatype_size_byte=4,
systolic_array_height=256,
systolic_array_width=256,
activation_fifo_depth=8,
accumulator_array_height=4096,
log_file_output_dir='.',
model_name='unnamed'):
"""
Similar to `tf.layers.Conv2D`, but with some differences:
1. Default kernel initializer is variance_scaling_initializer(2.0).
2. Default padding is 'same'.
3. Support 'split' argument to do group convolution.
Variable Names:
* ``W``: weights
* ``b``: bias
"""
if kernel_initializer is None:
if get_tf_version_tuple() <= (1, 12):
kernel_initializer = tf.contrib.layers.variance_scaling_initializer(2.0)
else:
kernel_initializer = tf.keras.initializers.VarianceScaling(2.0, distribution='untruncated_normal')
dilation_rate = shape2d(dilation_rate)
# group conv implementation
data_format = get_data_format(data_format, keras_mode=False)
in_shape = inputs.get_shape().as_list()
channel_axis = 3 if data_format == 'NHWC' else 1
in_channel = in_shape[channel_axis]
assert in_channel is not None, "[mpusim_conv2d] Input cannot have unknown channel!"
assert in_channel % split == 0
assert kernel_regularizer is None and bias_regularizer is None and activity_regularizer is None, \
"Not supported by group conv or dilated conv!"
out_channel = filters
assert out_channel % split == 0
assert dilation_rate == [1, 1] or get_tf_version_tuple() >= (1, 5), 'TF>=1.5 required for dilated conv.'
kernel_shape = shape2d(kernel_size)
filter_shape = kernel_shape + [in_channel / split, out_channel]
stride = shape4d(strides, data_format=data_format)
kwargs = dict(data_format=data_format)
if get_tf_version_tuple() >= (1, 5):
kwargs['dilations'] = shape4d(dilation_rate, data_format=data_format)
W = tf.get_variable(
'W', filter_shape, initializer=kernel_initializer)
if use_bias:
b = tf.get_variable('b', [out_channel], initializer=bias_initializer)
if split == 1:
conv = mpu_sim_conv2d_lib.mpu_sim_conv2d(inputs,
W,
activations_datatype_size_byte,
weights_datatype_size_byte,
results_datatype_size_byte,
systolic_array_height,
systolic_array_width,
activation_fifo_depth,
accumulator_array_height,
log_file_output_dir,
model_name,
stride,
padding.upper(),
**kwargs)
else:
inputs = tf.split(inputs, split, channel_axis)
kernels = tf.split(W, split, 3)
outputs = [mpu_sim_conv2d_lib.mpu_sim_conv2d(input_block,
kernel_block,
activations_datatype_size_byte,
weights_datatype_size_byte,
results_datatype_size_byte,
systolic_array_height,
systolic_array_width,
activation_fifo_depth,
accumulator_array_height,
log_file_output_dir,
model_name,
stride,
padding.upper(),
**kwargs)
for input_block, kernel_block in zip(inputs, kernels)]
conv = tf.concat(outputs, channel_axis)
ret = tf.nn.bias_add(conv, b, data_format=data_format) if use_bias else conv
if activation is not None:
ret = activation(ret)
ret = tf.identity(ret, name='output')
ret.variables = VariableHolder(W=W)
if use_bias:
ret.variables.b=b
return ret