def model(self):
d, s, m, r = self.model_params
# Feature Extraction
size = self.padding + 1
weights = tf.get_variable('w1', shape=[size, size, 1, d], initializer=tf.variance_scaling_initializer(0.1))
biases = tf.get_variable('b1', initializer=tf.zeros([d]))
features = tf.nn.conv2d(self.images, weights, strides=[1,1,1,1], padding='VALID', data_format='NHWC')
features = tf.nn.bias_add(features, biases, data_format='NHWC')
# Shrinking
if self.model_params[1] > 0:
features = self.prelu(features, 1)
weights = tf.get_variable('w2', shape=[1, 1, d, s], initializer=tf.variance_scaling_initializer(2))
biases = tf.get_variable('b2', initializer=tf.zeros([s]))
features = tf.nn.conv2d(features, weights, strides=[1,1,1,1], padding='SAME', data_format='NHWC')
features = tf.nn.bias_add(features, biases, data_format='NHWC')
else:
s = d
conv = features
# Mapping (# mapping layers = m)
with tf.variable_scope("mapping_block") as scope:
for ri in range(r):
for i in range(3, m + 3):
weights = tf.get_variable('w{}'.format(i), shape=[3, 3, s, s], initializer=tf.variance_scaling_initializer(2))
biases = tf.get_variable('b{}'.format(i), initializer=tf.zeros([s]))
if i > 3:
conv = self.prelu(conv, i)
conv = tf.nn.conv2d(conv, weights, strides=[1,1,1,1], padding='SAME', data_format='NHWC')
conv = tf.nn.bias_add(conv, biases, data_format='NHWC')
if i == m + 2:
conv = self.prelu(conv, m + 3)
weights = tf.get_variable('w{}'.format(m + 3), shape=[1, 1, s, s], initializer=tf.variance_scaling_initializer(2))
biases = tf.get_variable('b{}'.format(m + 3), initializer=tf.zeros([s]))
conv = tf.nn.conv2d(conv, weights, strides=[1,1,1,1], padding='SAME', data_format='NHWC')
conv = tf.nn.bias_add(conv, biases, data_format='NHWC')
conv = tf.add(conv, features)
scope.reuse_variables()
conv = self.prelu(conv, 2)
# Expanding
if self.model_params[1] > 0:
expand_weights = tf.get_variable('w{}'.format(m + 4), shape=[1, 1, s, d], initializer=tf.variance_scaling_initializer(2))
expand_biases = tf.get_variable('b{}'.format(m + 4), initializer=tf.zeros([d]))
conv = tf.nn.conv2d(conv, expand_weights, strides=[1,1,1,1], padding='SAME', data_format='NHWC')
conv = tf.nn.bias_add(conv, expand_biases, data_format='NHWC')
conv = self.prelu(conv, m + 4)
# Sub-pixel convolution
size = self.radius * 2 + 1
deconv_weights = tf.get_variable('deconv_w', shape=[size, size, d, self.scale**2], initializer=tf.variance_scaling_initializer(0.01))
deconv_biases = tf.get_variable('deconv_b', initializer=tf.zeros([self.scale**2]))
deconv = tf.nn.conv2d(conv, deconv_weights, strides=[1,1,1,1], padding='SAME', data_format='NHWC')
deconv = tf.nn.bias_add(deconv, deconv_biases, data_format='NHWC')
if self.scale > 1:
deconv = tf.depth_to_space(deconv, self.scale, name='pixel_shuffle', data_format='NHWC')
return deconv
def sk_conv2d(inputs, filters, strides, sk_ratio, min_dim=32,
is_training=True, data_format='channels_last'):
"""Selective kernel convolutional layer (https://arxiv.org/abs/1903.06586)."""
channel_axis = 1 if data_format == 'channels_first' else 3
pooling_axes = [2, 3] if data_format == 'channels_first' else [1, 2]
# Two stream convs (using split and both are 3x3).
inputs = conv2d_fixed_padding(
inputs=inputs, filters=2 * filters, kernel_size=3, strides=strides,
data_format=data_format)
inputs = batch_norm_relu(inputs, is_training, data_format=data_format)
inputs = tf.stack(tf.split(inputs, num_or_size_splits=2, axis=channel_axis))
# Mixing weights for two streams.
mid_dim = max(int(filters * sk_ratio), min_dim)
global_features = tf.reduce_mean(
tf.reduce_sum(inputs, axis=0), pooling_axes, keepdims=True)
global_features = tf.layers.conv2d(
inputs=global_features, filters=mid_dim, kernel_size=1, strides=1,
kernel_initializer=tf.variance_scaling_initializer(),
use_bias=False, data_format=data_format)
global_features = batch_norm_relu(
global_features, is_training, data_format=data_format)
mixing = tf.layers.conv2d(
inputs=global_features, filters=2 * filters, kernel_size=1, strides=1,
kernel_initializer=tf.variance_scaling_initializer(),
use_bias=False, data_format=data_format)
mixing = tf.stack(tf.split(mixing, num_or_size_splits=2, axis=channel_axis))
mixing = tf.nn.softmax(mixing, axis=0)
return tf.reduce_sum(inputs * mixing, axis=0)
def se_layer(inputs, filters, se_ratio, data_format='channels_last'):
"""Squeeze and Excitation layer (https://arxiv.org/abs/1709.01507)."""
if se_ratio <= 0:
return inputs
se_reduce = tf.layers.Conv2D(
max(1, int(filters * se_ratio)),
kernel_size=[1, 1],
strides=[1, 1],
kernel_initializer=tf.variance_scaling_initializer(),
padding='same',
data_format=data_format,
use_bias=True)
se_expand = tf.layers.Conv2D(
inputs.shape[-1],
kernel_size=[1, 1],
strides=[1, 1],
kernel_initializer=tf.variance_scaling_initializer(),
padding='same',
data_format=data_format,
use_bias=True)
spatial_dims = [2, 3] if data_format == 'channels_first' else [1, 2]
se_tensor = tf.reduce_mean(
inputs, spatial_dims, keepdims=True)
se_tensor = se_expand(tf.nn.relu(se_reduce(se_tensor)))
return tf.sigmoid(se_tensor) * inputs
def get_variable_initializer(hparams):
"""Get variable initializer from hparams."""
if not hparams.initializer:
return None
mlperf_log.transformer_print(key=mlperf_log.MODEL_HP_INITIALIZER_GAIN,
value=hparams.initializer_gain,
hparams=hparams)
if not tf.executing_eagerly():
tf.logging.info("Using variable initializer: %s", hparams.initializer)
if hparams.initializer == "orthogonal":
return tf.orthogonal_initializer(gain=hparams.initializer_gain)
elif hparams.initializer == "uniform":
max_val = 0.1 * hparams.initializer_gain
return tf.random_uniform_initializer(-max_val, max_val)
elif hparams.initializer == "normal_unit_scaling":
return tf.variance_scaling_initializer(
hparams.initializer_gain, mode="fan_avg", distribution="normal")
elif hparams.initializer == "uniform_unit_scaling":
return tf.variance_scaling_initializer(
hparams.initializer_gain, mode="fan_avg", distribution="uniform")
elif hparams.initializer == "xavier":
return tf.initializers.glorot_uniform()
else:
raise ValueError("Unrecognized initializer: %s" % hparams.initializer)
def __init__(self,
num_inputs,
num_outputs,
data_dim,
activation=tf.nn.relu,
name="resmade_block"):
self.num_inputs = num_inputs
self.num_outputs = num_outputs
self.activation = activation
with tf.variable_scope(name):
self.layer1 = MaskedDenseLayer(
num_inputs,
num_outputs,
data_dim,
activation=activation,
mask_type="hidden",
weight_initializer=tf.variance_scaling_initializer(
scale=2.0, distribution="normal"),
name="layer1")
self.layer2 = MaskedDenseLayer(
num_inputs,
num_outputs,
data_dim,
activation=None,
mask_type="hidden",
weight_initializer=tf.variance_scaling_initializer(
scale=0.1, distribution="normal"),
name="layer2")
def build_network(self):
with tf.compat.v1.variable_scope(self.name):
self.input = tf.placeholder(tf.float32,
shape=[None, *self.input_dims],
name='inputs')
self.actions = tf.placeholder(tf.float32,
shape=[None, self.n_actions],
name='action_taken')
self.q_target = tf.placeholder(tf.float32,
shape=[None, self.n_actions],
name='q_value')
conv1 = tf.layers.conv2d(
inputs=self.input,
filters=32,
kernel_size=(8, 8),
strides=4,
name='conv1',
kernel_initializer=tf.variance_scaling_initializer(scale=2))
conv1_activated = tf.nn.relu(conv1)
conv2 = tf.layers.conv2d(
inputs=conv1_activated,
filters=64,
kernel_size=(4, 4),
strides=2,
name='conv2',
kernel_initializer=tf.variance_scaling_initializer(scale=2))
conv2_activated = tf.nn.relu(conv2)
conv3 = tf.layers.conv2d(
inputs=conv2_activated,
filters=128,
kernel_size=(3, 3),
strides=1,
name='conv3',
kernel_initializer=tf.variance_scaling_initializer(scale=2))
conv3_activated = tf.nn.relu(conv3)
flat = tf.layers.flatten(conv3_activated)
dense1 = tf.layers.dense(
flat,
units=self.fc1_dims,
activation=tf.nn.relu,
kernel_initializer=tf.variance_scaling_initializer(scale=2))
self.Q_values = tf.layers.dense(
dense1,
units=self.n_actions,
kernel_initializer=tf.variance_scaling_initializer(scale=2))
#self.q = tf.reduce_sum(tf.multiply(self.Q_values, self.actions))
self.loss = tf.reduce_mean(tf.square(self.Q_values -
self.q_target))
self.train_op = tf.train.AdamOptimizer(self.lr).minimize(self.loss)
def model(self):
d = self.model_params
m = len(d) + 2
# Feature Extraction
size = self.padding + 1
weights = tf.get_variable(
'w1',
shape=[size, size, 1, d[0]],
initializer=tf.variance_scaling_initializer(0.1))
biases = tf.get_variable('b1', initializer=tf.zeros([d[0]]))
conv = tf.nn.conv2d(self.images,
weights,
strides=[1, 1, 1, 1],
padding='VALID',
data_format='NHWC')
conv = tf.nn.bias_add(conv, biases, data_format='NHWC')
conv = self.prelu(conv, 1)
# Mapping (# mapping layers = m)
for i in range(3, m):
weights = tf.get_variable(
'w{}'.format(i),
shape=[3, 3, d[i - 3], d[i - 2]],
initializer=tf.variance_scaling_initializer(2))
biases = tf.get_variable('b{}'.format(i),
initializer=tf.zeros([d[i - 2]]))
conv = tf.nn.conv2d(conv,
weights,
strides=[1, 1, 1, 1],
padding='SAME',
data_format='NHWC')
conv = tf.nn.bias_add(conv, biases, data_format='NHWC')
conv = self.prelu(conv, i)
# Sub-pixel convolution
size = self.radius * 2 + 1
deconv_weights = tf.get_variable(
'deconv_w',
shape=[size, size, d[-1], self.scale**2],
initializer=tf.variance_scaling_initializer(0.01))
deconv_biases = tf.get_variable('deconv_b',
initializer=tf.zeros([self.scale**2]))
deconv = tf.nn.conv2d(conv,
deconv_weights,
strides=[1, 1, 1, 1],
padding='SAME',
data_format='NHWC')
deconv = tf.nn.bias_add(deconv, deconv_biases, data_format='NHWC')
deconv = tf.depth_to_space(deconv,
self.scale,
name='pixel_shuffle',
data_format='NHWC')
return deconv
def __call__(self, input, **kwargs):
"""Call separable_conv2d function."""
with tf.variable_scope(self._scope_name, reuse=tf.AUTO_REUSE):
return tf.keras.layers.SeparableConv2D(filters=self.out_channels,
kernel_size=self.kernel_size,
strides=self.stride,
data_format=self.data_format,
dilation_rate=self.dilation,
depthwise_initializer=tf.variance_scaling_initializer(),
pointwise_initializer=tf.variance_scaling_initializer(),
padding='SAME', use_bias=self.bias,
name='SeparableConv2d',
reuse=self.reuse)(inputs=input)
def get_initializer(initializer, dtype):
if initializer == 'zeros':
return tf.zeros_initializer(dtype=dtype)
elif initializer == 'ones':
return tf.ones_initializer(dtype=dtype)
elif initializer == 'vs':
return tf.variance_scaling_initializer(dtype=dtype)
elif initializer == 'xavier':
return tf.glorot_normal_initializer(dtype=dtype)
elif initializer == 'he':
return tf.variance_scaling_initializer(dtype=dtype)
else:
raise NotImplementedError
def call(self, input, **kwargs):
"""Call separable_conv2d function."""
model = tf.keras.layers.SeparableConv2D(filters=self.out_channels,
kernel_size=self.kernel_size,
strides=self.stride,
data_format=self.data_format,
dilation_rate=self.dilation,
depthwise_initializer=tf.variance_scaling_initializer(),
pointwise_initializer=tf.variance_scaling_initializer(),
padding='SAME', use_bias=self.bias,
name=self.name,
reuse=self.reuse, trainable=self._trainable)
return model(inputs=input)
def squeeze_excitation(inputs,
in_filters,
se_ratio,
name=None,
expand_ratio=1,
data_format='channels_last'):
"""Squeeze and excitation implementation.
Args:
inputs: `Tensor` of size `[batch, channels, height_in, width_in]`.
in_filters: `int` number of input filteres before expansion.
se_ratio: `float` a se ratio between 0 and 1 for squeeze and excitation.
expand_ratio: `int` expansion ratio for the block.
data_format: An optional string from: "channels_last", "channels_first".
Defaults to "channels_last".
Returns:
A `Tensor` of shape `[batch, filters, height_out, width_out]`.
"""
num_reduced_filters = max(1, int(in_filters * se_ratio))
se_reduce = tf.layers.Conv2D(
num_reduced_filters,
kernel_size=[1, 1],
strides=[1, 1],
kernel_initializer=tf.variance_scaling_initializer(),
padding='same',
data_format=data_format,
name="se0_{}".format(name),
use_bias=True)
se_expand = tf.layers.Conv2D(
in_filters * expand_ratio,
kernel_size=[1, 1],
strides=[1, 1],
kernel_initializer=tf.variance_scaling_initializer(),
padding='same',
data_format=data_format,
name="se1_{}".format(name),
use_bias=True)
# Process input
if data_format == 'channels_first':
spatial_dims = [2, 3]
else:
spatial_dims = [1, 2]
se_tensor = tf.reduce_mean(inputs, spatial_dims, keepdims=True)
se_tensor = se_expand(tf.nn.relu(se_reduce(se_tensor)))
return tf.sigmoid(se_tensor) * inputs
def __init__(self, scope=None):
self.weights={}
self.scope=scope
self.kernel_initializer=tf.variance_scaling_initializer()
self.build_CNN_params()
print('Initialize weights {}'.format(self.scope))
def get_initializer(initializer, dtype):
if initializer == 'zeros':
return tf.zeros_initializer(dtype=dtype)
elif initializer == 'vs':
return tf.variance_scaling_initializer(dtype=dtype) #初始化
else:
raise NotImplementedError
def conv_linear_map(inputs, nin, nout, bias_start, prefix):
"""Convolutional liner map.
Maps 3D tensor by last dimension.
Args:
inputs: Inputs that should be shuffled
nin: Input feature map count
nout: Output feature map count
bias_start: Bias start value
prefix: Name prefix
Returns:
tf.Tensor: Inputs with applied convolution
"""
with tf.variable_scope(prefix):
inp_shape = tf.shape(inputs)
initializer = tf.variance_scaling_initializer(
scale=1.0, mode="fan_avg", distribution="uniform")
kernel = tf.get_variable("CvK", [nin, nout], initializer=initializer)
bias_term = tf.get_variable(
"CvB", [nout], initializer=tf.constant_initializer(0.0))
mul_shape = [inp_shape[0] * inp_shape[1], nin]
res = tf.matmul(tf.reshape(inputs, mul_shape), kernel)
res = tf.reshape(res, [inp_shape[0], inp_shape[1], nout])
return res + bias_start + bias_term
def create_conv_layer(filter_size, stride, input_size, output_size, input):
"""
:param filter_size: size of the filter
:param stride: stride
:param input_size: input size
:param output_size: number of neurons
:param input: input of the layer
:return A_conv: ouput
"""
with tf.variable_scope(None, default_name="conv"):
# Initialize the weights with tf.variance_scaling_initializer
W_init = tf.variance_scaling_initializer()
W_conv = tf.Variable(
W_init([filter_size, filter_size, input_size, output_size]))
# Initialize the biases with tf.zeros initializer
b_init = tf.zeros_initializer()
b_conv = tf.Variable(b_init(output_size))
A_conv = tf.nn.conv2d(
input, W_conv, strides=[1, stride, stride, 1
], padding='SAME') + b_conv
# Apply activation function
A_conv = tf.nn.relu(A_conv)
return A_conv
def prn(x, is_training):
"""
Arguments:
x: a float tensor with shape [b, h, w, c].
is_training: a boolean.
Returns:
a float tensor with shape [b, h, w, c].
"""
with tf.variable_scope('PRN'):
b = tf.shape(x)[0]
_, h, w, c = x.shape.as_list() # must be static
x = tf.reshape(x, [b, h * w * c]) # flatten
with slim.arg_scope(
[slim.fully_connected],
weights_initializer=tf.variance_scaling_initializer()):
y = slim.fully_connected(x,
1024,
activation_fn=tf.nn.relu,
scope='fc1')
# y = slim.dropout(y, keep_prob=0.5, is_training=is_training)
y = slim.fully_connected(y,
h * w * c,
activation_fn=tf.nn.relu,
scope='fc2')
x += y
return tf.reshape(x, [b, h, w, c])
def conv2d_fixed_padding(inputs,
filters,
kernel_size,
strides,
data_format='channels_first'):
"""Strided 2-D convolution with explicit padding.
The padding is consistent and is based only on `kernel_size`, not on the
dimensions of `inputs` (as opposed to using `tf.layers.conv2d` alone).
Args:
inputs: `Tensor` of size `[batch, channels, height_in, width_in]`.
filters: `int` number of filters in the convolution.
kernel_size: `int` size of the kernel to be used in the convolution.
strides: `int` strides of the convolution.
data_format: `str` either "channels_first" for `[batch, channels, height,
width]` or "channels_last for `[batch, height, width, channels]`.
Returns:
A `Tensor` of shape `[batch, filters, height_out, width_out]`.
"""
if strides > 1:
inputs = fixed_padding(inputs, kernel_size, data_format=data_format)
return tf.layers.conv2d(
inputs=inputs,
filters=filters,
kernel_size=kernel_size,
strides=strides,
padding=('SAME' if strides == 1 else 'VALID'),
use_bias=False,
kernel_initializer=tf.variance_scaling_initializer(),
data_format=data_format)
def conv_bn_relu(inputs, conv_size, conv_filters, is_training, data_format):
"""Convolution followed by batch norm and ReLU."""
if data_format == 'channels_last':
axis = 3
elif data_format == 'channels_first':
axis = 1
else:
raise ValueError('invalid data_format')
net = tf.layers.conv2d(
inputs=inputs,
filters=conv_filters,
kernel_size=conv_size,
strides=(1, 1),
use_bias=False,
kernel_initializer=tf.variance_scaling_initializer(),
padding='same',
data_format=data_format)
net = tf.layers.batch_normalization(inputs=net,
axis=axis,
momentum=BN_MOMENTUM,
epsilon=BN_EPSILON,
training=is_training)
net = tf.nn.relu(net)
return net
def se_block(input_feature, name, ratio=8):
kernel_initializer = tf.variance_scaling_initializer()
bias_initializer = tf.constant_initializer(value=0.0)
with tf.variable_scope(name):
channel = input_feature.get_shape()[-1]
# Global average pooling
squeeze = tf.reduce_mean(input_feature, axis=[1,2], keepdims=True)
assert squeeze.get_shape()[1:] == (1,1,channel)
excitation = tf.layers.dense(inputs=squeeze,
units=channel//ratio,
activation=tf.nn.relu,
kernel_initializer=kernel_initializer,
bias_initializer=bias_initializer,
name='bottleneck_fc', trainable=False)
assert excitation.get_shape()[1:] == (1,1,channel//ratio)
excitation = tf.layers.dense(inputs=excitation,
units=channel,
activation=tf.nn.sigmoid,
kernel_initializer=kernel_initializer,
bias_initializer=bias_initializer,
name='recover_fc', trainable=False)
assert excitation.get_shape()[1:] == (1,1,channel)
scale = input_feature * excitation
return scale
def build(self, input_shape):
"""Initialize layer weights and sublayers.
Args:
input_shape: shape of inputs
"""
in_units = input_shape[-1]
middle_units = in_units * 4
out_units = in_units * 2
init = tf.variance_scaling_initializer(scale=1.0,
mode="fan_avg",
distribution="uniform")
self.first_linear = tf.keras.layers.Dense(middle_units,
use_bias=False,
kernel_initializer=init,
name=self.prefix + "/cand1")
self.second_linear = tf.keras.layers.Dense(out_units,
kernel_initializer=init,
name=self.prefix + "/cand2")
self.layer_norm = LayerNormalization()
init = tf.constant_initializer(self.init_value)
self.residual_scale = self.add_weight(self.prefix + "/residual",
[out_units],
initializer=init)
super(RSU, self).build(input_shape)
def __init__(self,
in_channels,
out_channels,
kernel_size=3,
stride=1,
padding=0,
bias=True,
groups=1,
dilation=1,
separable=False,
depthwise=False,
padding_mode='same',
bn=False):
super(Conv2d, self).__init__()
self.in_channels = in_channels
self.out_channels = out_channels
self.kernel_size = kernel_size
self.stride = stride
self.padding = padding
self.bias = bias
self.groups = groups
self.dilation = dilation
self.kernel_initial = tf.variance_scaling_initializer()
self.bias_initial = tf.zeros_initializer()
self._initializer = None
self.reuse = None
self.separable = separable
self.depthwise = depthwise
self.padding_mode = padding if isinstance(
padding, str) and stride != 2 else padding_mode
self.bn = bn
def initialize_parameters():
tf.set_random_seed(1)
W1 = tf.get_variable("W1", [25, 12288],
initializer=tf.variance_scaling_initializer(seed=1))
b1 = tf.get_variable("b1", [25, 1], initializer=tf.zeros_initializer())
W2 = tf.get_variable("W2", [12, 25],
initializer=tf.variance_scaling_initializer(seed=1))
b2 = tf.get_variable("b2", [12, 1], initializer=tf.zeros_initializer())
W3 = tf.get_variable("W3", [6, 12],
initializer=tf.variance_scaling_initializer(seed=1))
b3 = tf.get_variable("b3", [6, 1], initializer=tf.zeros_initializer())
parameters = {"W1": W1, "b1": b1, "W2": W2, "b2": b2, "W3": W3, "b3": b3}
return parameters
def _network_template(self, obs, num_layers, hidden_units):
"""PixelCNN network architecture."""
with slim.arg_scope(
[slim.conv2d, masked_conv2d],
weights_initializer=tf.variance_scaling_initializer(
distribution='uniform'),
biases_initializer=tf.constant_initializer(0.0)):
net = masked_conv2d(obs,
hidden_units, [7, 7],
mask_type='A',
activation_fn=None,
scope='masked_conv_1')
embedding = slim.model_variable(
'embedding',
shape=(1, ) + self.resize_shape + (4, ),
initializer=tf.variance_scaling_initializer(
distribution='uniform'))
for i in range(1, num_layers + 1):
net2 = gating_layer(net, embedding, hidden_units,
'gating_{}'.format(i))
net += masked_conv2d(net2,
hidden_units, [1, 1],
mask_type='B',
activation_fn=None,
scope='masked_conv_{}'.format(i + 1))
net += slim.conv2d(embedding,
hidden_units, [1, 1],
activation_fn=None)
net = tf.nn.relu(net)
net = masked_conv2d(net,
64, [1, 1],
scope='1x1_conv_out',
mask_type='B',
activation_fn=tf.nn.relu)
logits = masked_conv2d(net,
self.quantization_factor, [1, 1],
scope='logits',
mask_type='B',
activation_fn=None)
loss = tf.losses.sparse_softmax_cross_entropy(
labels=tf.cast(obs, tf.int32),
logits=logits,
reduction=tf.losses.Reduction.MEAN)
return collections.namedtuple('PixelCNN_network',
['logits', 'loss'])(logits, loss)
def conv0_space_to_depth(inputs, filters, kernel_size, strides,
data_format='channels_last',
space_to_depth_block_size=2):
"""Uses space-to-depth convolution for conv0.
This function replaces the first convolution (conv0) in ResNet with
space-to-depth transformation. It creates a convolution kernel, whose
dimension and name are the same as those of conv0. The `inputs` is an image
tensor that already has the space-to-depth transform.
Args:
inputs: `Tensor` of size `[batch, height_in, width_in, channels]`.
filters: An `int` number of filters in the convolution.
kernel_size: An `int` size of the kernel to be used in the convolution.
strides: A `int` strides of the convolution.
data_format: A `str` either "channels_first" for `[batch, channels, height,
width]` or "channels_last for `[batch, height, width, channels]`.
space_to_depth_block_size: An `int` indicates the block size of
space-to-depth convolution for conv0. Specific to ResNet, this currently
supports only block_size=2.
Returns:
A `Tensor` with the same type as `inputs`.
Raises:
ValueError if `space_to_depth_block_size` is not 2.
"""
if space_to_depth_block_size != 2:
raise ValueError('Space-to-depth does not support block_size (%d).' %
space_to_depth_block_size)
conv0 = tf.layers.Conv2D(
filters=filters,
kernel_size=kernel_size,
strides=strides,
padding=('SAME' if strides == 1 else 'VALID'),
data_format=data_format,
use_bias=False,
kernel_initializer=tf.variance_scaling_initializer())
# Use the image size without space-to-depth transform as the input of conv0.
# This allows the kernel size to be the same as the original conv0 such that
# the model is able to load the pre-trained ResNet checkpoint.
batch_size, h, w, c = inputs.get_shape().as_list()
conv0.build([batch_size,
h * space_to_depth_block_size,
w * space_to_depth_block_size,
c // (space_to_depth_block_size ** 2)])
kernel = conv0.weights[0]
kernel = transform_space_to_depth_kernel(
kernel, inputs.dtype, block_size=space_to_depth_block_size)
inputs = spatial_transform.space_to_depth_fixed_padding(
inputs, kernel_size, data_format, space_to_depth_block_size)
return tf.nn.conv2d(
input=inputs, filter=kernel, strides=[1, 1, 1, 1], padding='VALID',
data_format='NHWC' if data_format == 'channels_last' else 'NCHW',
name='conv2d/Conv2D')
def conv2d_fixed_padding(inputs,
filters,
kernel_size,
strides,
pruning_method='baseline',
data_format='channels_first',
weight_decay=0.,
name=None):
"""Strided 2-D convolution with explicit padding.
The padding is consistent and is based only on `kernel_size`, not on the
dimensions of `inputs` (as opposed to using `tf.layers.conv2d` alone).
Args:
inputs: Input tensor, float32 or bfloat16 of size [batch, channels, height,
width].
filters: Int specifying number of filters for the first two convolutions.
kernel_size: Int designating size of kernel to be used in the convolution.
strides: Int specifying the stride. If stride >1, the input is downsampled.
pruning_method: String that specifies the pruning method used to identify
which weights to remove.
data_format: String that specifies either "channels_first" for [batch,
channels, height,width] or "channels_last" for [batch, height, width,
channels].
weight_decay: Weight for the l2 regularization loss.
name: String that specifies name for model layer.
Returns:
The output activation tensor of size [batch, filters, height_out, width_out]
Raises:
ValueError: If the data_format provided is not a valid string.
"""
if strides > 1:
inputs = resnet_model.fixed_padding(inputs,
kernel_size,
data_format=data_format)
padding = 'VALID'
else:
padding = 'SAME'
kernel_initializer = tf.variance_scaling_initializer()
kernel_regularizer = contrib_layers.l2_regularizer(weight_decay)
return sparse_conv2d(x=inputs,
units=filters,
activation=None,
kernel_size=[kernel_size, kernel_size],
use_bias=False,
kernel_initializer=kernel_initializer,
kernel_regularizer=kernel_regularizer,
bias_initializer=None,
biases_regularizer=None,
sparsity_technique=pruning_method,
normalizer_fn=None,
strides=[strides, strides],
padding=padding,
data_format=data_format,
name=name)
def conv0_space_to_depth(inputs, data_format='channels_last'):
"""Strided 2-D convolution with explicit padding.
The padding is consistent and is based only on `kernel_size`, not on the
dimensions of `inputs` (as opposed to using `tf.layers.conv2d` alone).
Args:
inputs: `Tensor` of size `[batch, height_in, width_in, channels]`.
data_format: `str` either "channels_first" for `[batch, channels, height,
width]` or "channels_last for `[batch, height, width, channels]`.
Returns:
A `Tensor` with the same type as `inputs`.
"""
# Create the conv0 kernel w.r.t. the original image size. (no space-to-depth).
filters = 64
kernel_size = 7
space_to_depth_block_size = ssd_constants.SPACE_TO_DEPTH_BLOCK_SIZE
strides = 2
conv0 = tf.compat.v1.layers.Conv2D(
filters=filters,
kernel_size=kernel_size,
strides=2,
padding=('SAME' if strides == 1 else 'VALID'),
use_bias=False,
kernel_initializer=tf.variance_scaling_initializer(),
data_format=data_format)
# Use the image size without space-to-depth transform as the input of conv0.
batch_size, h, w, channel = inputs.get_shape().as_list()
conv0.build([
batch_size, h * space_to_depth_block_size,
w * space_to_depth_block_size,
channel // (space_to_depth_block_size**2)
])
kernel = conv0.weights[0]
# [7, 7, 3, 64] --> [8, 8, 3, 64]
kernel = tf.pad(kernel,
paddings=tf.constant([[1, 0], [1, 0], [0, 0], [0, 0]]),
mode='CONSTANT',
constant_values=0.)
# Transform kernel follows the space-to-depth logic: http://shortn/_9YvHW96xPJ
kernel = tf.reshape(kernel, [
4, space_to_depth_block_size, 4, space_to_depth_block_size, 3, filters
])
kernel = tf.transpose(kernel, [0, 2, 1, 3, 4, 5])
kernel = tf.reshape(kernel, [4, 4, int(channel), filters])
kernel = tf.cast(kernel, inputs.dtype)
inputs = space_to_depth_fixed_padding(inputs, kernel_size, data_format,
space_to_depth_block_size)
return tf.nn.conv2d(
input=inputs,
filter=kernel,
strides=[1, 1, 1, 1],
padding='VALID',
data_format='NHWC' if data_format == 'channels_last' else 'NCHW',
name='conv2d/Conv2D')
def __init__(self, sparsity, seed=None, dtype=tf.float32):
if sparsity < 0. or sparsity > 1.:
raise ValueError('sparsity must be in the range [0., 1.].')
self.kernel_initializer = tf.variance_scaling_initializer(seed=seed,
dtype=dtype)
self.seed = seed
self.dtype = dtype
self.sparsity = float(sparsity)
def conv2d_fixed_padding(inputs, filters, kernel_size, strides, data_format):
if strides > 1:
inputs = fixed_padding(inputs, kernel_size, data_format)
return tf.layers.conv2d(
inputs=inputs, filters=filters, kernel_size=kernel_size, strides=strides,
padding=('SAME' if strides == 1 else 'VALID'), use_bias=False,
kernel_initializer=tf.variance_scaling_initializer(),
data_format=data_format)
def conv2d_space_to_depth(inputs, filters, kernel_size, strides, block_size=2):
"""Strided 2-D convolution with explicit padding.
The padding is consistent and is based only on `kernel_size`, not on the
dimensions of `inputs` (as opposed to using `tf.layers.conv2d` alone).
Args:
inputs: `Tensor` of size `[batch, height_in, width_in, channels]`.
filters: `int` number of filters in the convolution.
kernel_size: `int` size of the kernel to be used in the convolution.
strides: `int` strides of the convolution.
block_size: space-to-depth padding block-size
Returns:
A `Tensor` of shape `[batch, filters, height_out, width_out]`.
"""
def _round_up(kernel_size, block_size):
remainder = kernel_size % block_size
if remainder == 0:
return block_size
else:
return kernel_size + block_size - remainder
padded_k = _round_up(kernel_size, block_size)
_, _, _, in_f = inputs.get_shape().as_list()
# Create the padded kernel
kernel = tf.get_variable(
'kernel_conv2d_opt',
shape=[padded_k, padded_k, in_f // 4, filters],
initializer=tf.variance_scaling_initializer(),
trainable=True,
dtype=tf.float32)
# Zero padded region
mx = tf.constant([1, 1, 1, 1, 1, 1, 1, 0], dtype=tf.float32)
kernel = kernel * tf.reshape(mx, [1, padded_k, 1, 1]) * tf.reshape(
mx, [padded_k, 1, 1, 1])
# Transpose to enable space-to-depth optimization
kernel = tf.reshape(kernel, [
padded_k // 2, block_size, padded_k // 2, block_size, in_f // 4, filters
])
kernel = tf.transpose(kernel, [0, 2, 1, 3, 4, 5])
kernel = tf.reshape(kernel, [padded_k // 2, padded_k // 2, in_f, filters])
kernel = tf.cast(kernel, inputs.dtype)
# Add a convolution operation using this kernel
s = strides // 2
return tf.nn.conv2d(
input=inputs,
filter=kernel,
strides=[s, s, s, s],
padding='SAME',
data_format='NHWC')
def default_mlp(hidden_sizes, activate_final=False, init_std=2., **kwargs):
"""Standard batch-applied MLP for transformer modules."""
init = {
'w': tf.variance_scaling_initializer(init_std, distribution='normal')
}
mlp = snt_mlp.MLP(hidden_sizes,
activate_final=activate_final,
use_dropout=True,
initializers=init,
**kwargs)
return basic.BatchApply(mlp)
