def __init__(self, kernel_size=4, name=None):
super(AutoEncoder, self).__init__(name=name)
self.encoder = snt.Sequential([
snt.Conv2D(4, kernel_size, stride=4, padding='SAME'), tf.nn.relu,
snt.Conv2D(16, kernel_size, stride=4, padding='SAME'), tf.nn.relu,
snt.Conv2D(64, kernel_size, stride=4, padding='SAME'), tf.nn.relu
])
self.decoder = snt.Sequential([
snt.Conv2DTranspose(64, kernel_size, stride=4,
padding='SAME'), tf.nn.relu,
snt.Conv2DTranspose(16, kernel_size, stride=4,
padding='SAME'), tf.nn.relu,
snt.Conv2DTranspose(4, kernel_size, stride=4, padding='SAME'),
tf.nn.relu,
snt.Conv2D(1, kernel_size, padding='SAME')
])
def __init__(self, action_spec):
super().__init__(name='r2d2_test_network')
self._net = snt.DeepRNN([
snt.Conv2D(32, [8, 8], [4, 4]),
tf.nn.relu,
snt.Conv2D(64, [4, 4], [2, 2]),
tf.nn.relu,
snt.Conv2D(64, [3, 3], [1, 1]),
tf.nn.relu,
snt.Flatten(),
snt.LSTM(20),
tf.nn.relu,
#snt.LSTM(160),
#snt.nets.MLP([50, 50,512]),
#tf.nn.relu,
snt.nets.MLP([50, 50, action_spec])
])
def testGetSaverModule(self):
input_ = tf.placeholder(tf.float32, shape=[1, 10, 10, 3])
conv = snt.Conv2D(output_channels=3, kernel_shape=3)
conv(input_)
saver = snt.get_saver(conv)
self.assertIsInstance(saver, tf.train.Saver)
self.assertIn("w:0", saver._var_list)
self.assertIn("b:0", saver._var_list)
def __init__(self, channel, kernel, name='unet_block_up'):
super(UNetBlockUp, self).__init__(name=name)
with self._enter_variable_scope(check_same_graph=False):
self._layers = [
snt.Conv2D(channel, kernel, use_bias=False, name='conv'),
snt.LayerNorm(axis=[1, 2], offset=True, scale=False, name='instance_norm'),
partial(tf.nn.relu, name='relu'),
]
def _build(self, z0, is_training=False):
"""Outputs logits."""
zk = z0
conv2d_id = 0
linear_id = 0
name = None
for spec in self._layer_types:
if spec[0] == 'conv2d':
if linear_id > 0:
raise ValueError(
'Convolutional layers must precede fully connected '
'layers.')
name = 'conv2d_{}'.format(conv2d_id)
conv2d_id += 1
(_, (kernel_height, kernel_width), channels, padding,
stride) = spec
m = snt.Conv2D(output_channels=channels,
kernel_shape=(kernel_height, kernel_width),
padding=padding,
stride=stride,
use_bias=True,
regularizers=self._regularizers,
initializers=_create_conv2d_initializer(
zk.get_shape().as_list()[1:], channels,
(kernel_height, kernel_width)),
name=name)
zk = m(zk)
elif spec[0] == 'linear':
must_flatten = (linear_id == 0 and len(zk.shape) > 2)
if must_flatten:
zk = snt.BatchFlatten()(zk)
name = 'linear_{}'.format(linear_id)
linear_id += 1
output_size = spec[1]
m = snt.Linear(output_size,
regularizers=self._regularizers,
initializers=_create_linear_initializer(
np.prod(zk.get_shape().as_list()[1:]),
output_size),
name=name)
zk = m(zk)
elif spec[0] == 'batch_normalization':
if name is None:
raise ValueError(
'Batch normalization only supported after linear '
'layers.')
name += '_batch_norm'
m = layers.BatchNorm(name=name)
zk = m(zk, is_training=is_training)
elif spec[0] == 'activation':
if spec[1] not in _ALLOWED_ACTIVATIONS:
raise NotImplementedError(
'Only the following activations are supported {}'.
format(list(_ALLOWED_ACTIVATIONS)))
name = None
m = getattr(tf.nn, spec[1])
zk = m(zk)
return zk
def __init__(self, kernel_size=4, name=None):
super(AutoEncoder, self).__init__(name=name)
self.encoder = snt.Sequential([snt.Conv2D(4, kernel_size, stride=2, padding='SAME'), tf.nn.leaky_relu, # [4, 128, 128]
snt.Conv2D(8, kernel_size, stride=2, padding='SAME'), tf.nn.leaky_relu, # [8, 64, 64]
snt.Conv2D(16, kernel_size, stride=2, padding='SAME'), tf.nn.leaky_relu, # [16, 32, 32]
snt.Conv2D(32, kernel_size, stride=2, padding='SAME'), tf.nn.leaky_relu]) # [32, 16, 16]
# snt.Conv2D(32, kernel_size, stride=2, padding='SAME'), tf.nn.relu,
# snt.Conv2D(64, kernel_size, stride=2, padding='SAME'), tf.nn.relu])
# self.decoder = snt.Sequential([snt.Conv2DTranspose(64, kernel_size, stride=2, padding='SAME'), tf.nn.relu,
# snt.Conv2DTranspose(32, kernel_size, stride=2, padding='SAME'), tf.nn.relu,
# snt.Conv2DTranspose(16, kernel_size, stride=2, padding='SAME'), tf.nn.relu,
self.decoder = snt.Sequential([snt.Conv2DTranspose(32, kernel_size, stride=2, padding='SAME'), tf.nn.leaky_relu, # [32, 16, 16]
snt.Conv2DTranspose(16, kernel_size, stride=2, padding='SAME'), tf.nn.leaky_relu, # [16, 32, 32]
snt.Conv2DTranspose(8, kernel_size, stride=2, padding='SAME'), tf.nn.leaky_relu, # [8, 64, 64]
snt.Conv2DTranspose(4, kernel_size, stride=2, padding='SAME'), tf.nn.leaky_relu, # [4, 128, 128]
snt.Conv2D(1, kernel_size, padding='SAME')]) # [1, 256, 256]
def residual_stack(h, num_hiddens, num_residual_layers, num_residual_hiddens):
for i in range(num_residual_layers):
h_i = tf.nn.relu(h)
h_i = snt.Conv2D(
output_channels=num_residual_hiddens,
kernel_shape=(3, 3),
stride=(1, 1),
name="res3x3_%d" % i)(h_i)
h_i = tf.nn.relu(h_i)
h_i = snt.Conv2D(
output_channels=num_hiddens,
kernel_shape=(1, 1),
stride=(1, 1),
name="res1x1_%d" % i)(h_i)
h += h_i
return tf.nn.relu(h)
def test_count_parameters_on_module(self):
module = snt.Module()
# Weights of a 2D convolution with 2 filters..
module.conv = snt.Conv2D(output_channels=2,
kernel_shape=3,
name="conv")
module.conv(tf.ones(
(2, 5, 5, 3))) # 3 * 3*3 * 2 + 2 (bias) = 56 parameters
self.assertEqual(56, parameter_overview.count_parameters(module))
def _build(self, x):
h = x
for unused_i, l in enumerate(self.layers):
h = tf.nn.relu(snt.Conv2D(l[0], l[1], l[2])(h))
h_shape = h.get_shape().as_list()
h = tf.reshape(h, [-1, h_shape[1] * h_shape[2] * h_shape[3]])
logits = snt.Linear(self.output_size)(h)
return logits
def _torso(self, input_):
last_action, env_output = input_
reward, _, _, (frame, instruction) = env_output
# Convert to floats.
frame = tf.to_float(frame)
frame /= 255
with tf.variable_scope('convnet'):
conv_out = frame
for i, (num_ch, num_blocks) in enumerate([(16, 2), (32, 2), (32, 2)]):
# Downscale.
conv_out = snt.Conv2D(num_ch, 3, stride=1, padding='SAME')(conv_out)
conv_out = tf.nn.pool(
conv_out,
window_shape=[3, 3],
pooling_type='MAX',
padding='SAME',
strides=[2, 2])
# Residual block(s).
for j in range(num_blocks):
with tf.variable_scope('residual_%d_%d' % (i, j)):
block_input = conv_out
conv_out = tf.nn.relu(conv_out)
conv_out = snt.Conv2D(num_ch, 3, stride=1, padding='SAME')(conv_out)
conv_out = tf.nn.relu(conv_out)
conv_out = snt.Conv2D(num_ch, 3, stride=1, padding='SAME')(conv_out)
conv_out += block_input
conv_out = tf.nn.relu(conv_out)
conv_out = snt.BatchFlatten()(conv_out)
conv_out = snt.Linear(256)(conv_out)
conv_out = tf.nn.relu(conv_out)
instruction_out = self._instruction(instruction)
# Append clipped last reward and one hot last action.
clipped_reward = tf.expand_dims(tf.clip_by_value(reward, -1, 1), -1)
one_hot_last_action = tf.one_hot(last_action, self._num_actions)
return tf.concat(
[conv_out, clipped_reward, one_hot_last_action, instruction_out],
axis=1)
def __init__(self, filter_size = 3, num_filters = 32,
pooling_stride = 2, act = 'tanh', summ = None, name = "mapper"):
super(Mapper, self).__init__(name = name)
self._pool = Downsample2D(pooling_stride)
self._act = Activation(act, verbose = True)
self._bf = snt.BatchFlatten()
self._summ = summ
initializers = {
'w': tf.truncated_normal_initializer(stddev = 0.02),
'b': tf.zeros_initializer()
}
with self._enter_variable_scope():
self._l1_conv = snt.Conv2D(num_filters, filter_size)
self._l2_conv = snt.Conv2D(num_filters << 1, filter_size)
self._lin1 = snt.Linear(256, initializers = initializers)
self._lin2 = snt.Linear(1, initializers = initializers)
def _build(self, inputs, verbose=VERBOSITY):
if EncodeProcessDecode_v8.convnet_tanh:
activation = tf.nn.tanh
else:
activation = tf.nn.relu
img_shape = get_correct_image_shape(config=None, get_type="seg",
depth_data_provided=EncodeProcessDecode_v8.depth_data_provided)
img_data = tf.reshape(inputs, [-1, *img_shape]) # -1 means "all", i.e. batch dimension
print(img_data.get_shape())
''' 60, 80 '''
outputs = snt.Conv2D(output_channels=32, kernel_shape=3, stride=2, padding="SAME")(img_data)
outputs = activation(outputs)
if EncodeProcessDecode_v8.conv_layer_instance_norm:
outputs = snt.BatchNorm()(outputs, is_training=self._is_training)
print(outputs.get_shape())
''' 30, 40 '''
outputs = snt.Conv2D(output_channels=32, kernel_shape=3, stride=2, padding="SAME")(outputs)
outputs = activation(outputs)
if EncodeProcessDecode_v8.conv_layer_instance_norm:
outputs = snt.BatchNorm()(outputs, is_training=self._is_training)
print(outputs.get_shape())
''' 15, 20 '''
outputs = snt.Conv2D(output_channels=16, kernel_shape=3, stride=2, padding="SAME")(outputs)
outputs = activation(outputs)
if EncodeProcessDecode_v8.conv_layer_instance_norm:
outputs = snt.BatchNorm()(outputs, is_training=self._is_training)
print(outputs.get_shape())
''' 8, 10 '''
outputs = snt.Conv2D(output_channels=5, kernel_shape=3, stride=2, padding="SAME")(outputs)
outputs = activation(outputs)
if EncodeProcessDecode_v8.conv_layer_instance_norm:
outputs = snt.BatchNorm()(outputs, is_training=self._is_training)
print(outputs.get_shape())
outputs = tf.layers.flatten(outputs) # 8,10,5 flattened
return outputs
def _create_conv(self, partitioned, name):
hidden = tf.ones(shape=(1, 16, 16, 3))
if partitioned:
partitioners = {"w": tf.variable_axis_size_partitioner(4)}
else:
partitioners = None
conv = snt.Conv2D(output_channels=3, kernel_shape=3, stride=1,
partitioners=partitioners, name=name)
conv(hidden)
return conv
def _build_conv_layer(conv_spec, data_format):
return snt.Conv2D(output_channels=conv_spec.output_channels,
kernel_shape=conv_spec.kernel_shape,
stride=conv_spec.stride,
rate=conv_spec.rate,
padding=snt.SAME,
use_bias=True,
data_format=data_format,
initializers=_DEFAULT_CONV_INITIALIZERS,
regularizers=_DEFAULT_CONV_REGULARIZERS)
def inference(inputs):
inputs = snt.Conv2D(output_channels=166,
kernel_shape=3,
rate=1,
padding='SAME',
name='conv1')(inputs)
inputs = tf.nn.relu(inputs)
inputs = tf.nn.max_pool(inputs, [1, 2, 2, 1], [1, 2, 2, 1], padding='SAME')
inputs = snt.Conv2D(output_channels=32,
kernel_shape=3,
rate=2,
padding='SAME',
name='conv2')(inputs)
inputs = tf.nn.relu(inputs)
inputs = tf.nn.max_pool(inputs, [1, 2, 2, 1], [1, 2, 2, 1], padding='SAME')
inputs = snt.Conv2D(output_channels=64,
kernel_shape=3,
rate=4,
padding='SAME',
name='conv3')(inputs)
inputs = tf.nn.relu(inputs)
inputs = tf.nn.max_pool(inputs, [1, 2, 2, 1], [1, 2, 2, 1], padding='SAME')
inputs = snt.Conv2D(output_channels=128,
kernel_shape=3,
rate=8,
padding='SAME',
name='conv4')(inputs)
inputs = tf.nn.relu(inputs)
inputs = tf.nn.max_pool(inputs, [1, 2, 2, 1], [1, 2, 2, 1], padding='SAME')
inputs = snt.Conv2D(output_channels=256,
kernel_shape=3,
rate=16,
padding='SAME',
name='conv5')(inputs)
inputs = tf.nn.relu(inputs)
inputs = tf.nn.max_pool(inputs, [1, 2, 2, 1], [1, 2, 2, 1], padding='SAME')
inputs = snt.Conv2D(output_channels=256,
kernel_shape=3,
padding='SAME',
name='conv6')(inputs)
inputs = tf.nn.relu(inputs)
inputs = tf.nn.max_pool(inputs, [1, 2, 2, 1], [1, 2, 2, 1], padding='SAME')
inputs = snt.Conv2D(output_channels=1,
kernel_shape=1,
padding='SAME',
name='conv7')(inputs)
coords, norm_heatmap = dsnt.dsnt(inputs)
# The Sonnet option
# coords, norm_heatmap = DSNT()(inputs)
return coords, norm_heatmap
def res_net_convolution(frame):
for i, (num_ch, num_blocks) in enumerate([(16, 2), (32, 2), (32, 2)]):
# Downscale.
conv_out = snt.Conv2D(num_ch, 3, stride=1, padding='SAME')(frame)
conv_out = tf.nn.pool(
conv_out,
window_shape=[3, 3],
pooling_type='MAX',
padding='SAME',
strides=[2, 2])
# Residual block(s).
for j in range(num_blocks):
with tf.variable_scope('residual_%d_%d' % (i, j)):
block_input = conv_out
conv_out = tf.nn.relu(conv_out)
conv_out = snt.Conv2D(num_ch, 3, stride=1, padding='SAME')(conv_out)
conv_out = tf.nn.relu(conv_out)
conv_out = snt.Conv2D(num_ch, 3, stride=1, padding='SAME')(conv_out)
conv_out += block_input
return conv_out
def decoder_builder(y, z):
# TODO: remove the linear layer
# `n` is local parameter for size of internal linear layer
n = 32
latent = snt.Linear(n)(tf.concat(
[tf.cast(y, dtype=tf.float32),
tf.cast(z, dtype=tf.float32)], 1))
latent = tf.reshape(latent, [-1, 1, 1, n])
output = network(latent, **network_kwargs)
output = snt.Conv2D(num_channels, 1, padding=snt.SAME)(output)
return output
def test_leading_batchnorm_rejected(self):
module = snt.Sequential([
_BatchNorm(),
snt.Conv2D(output_channels=5, kernel_shape=3, padding='VALID'),
])
network = ibp.VerifiableModelWrapper(module)
network(self._inputs)
with self.assertRaises(auto_verifier.NotVerifiableError):
_ = auto_verifier.VerifiableLayerBuilder(network).build_layers()
def __init__(self,
image_feature_size=16,
kernel_size=3,
name=None,
**unused_kwargs):
super(Model, self).__init__(name=name)
self.encoder_graph = RelationNetwork(
lambda: snt.nets.MLP([32, 16], activate_final=True),
lambda: snt.nets.MLP([32, 16], activate_final=True))
self.encoder_image = RelationNetwork(
lambda: snt.nets.MLP([32, 16], activate_final=True),
lambda: snt.nets.MLP([32, 16], activate_final=True))
self.image_cnn = snt.Sequential([
snt.Conv2D(16, kernel_size, stride=2), tf.nn.relu,
snt.Conv2D(image_feature_size, kernel_size, stride=2), tf.nn.relu
])
self.compare = snt.nets.MLP([32, 1])
self.image_feature_size = image_feature_size
self._step = None
def _avgpool_linear_script(num_classes, net, layer_values):
layer_values = add_layer(net, snt.Conv2D(3, kernel_shape=(2, 2), **_inits),
layer_values)
layer_values = add_layer(net, tf.nn.relu, layer_values)
layer_values = add_layer(net, AvgPool(kernel_shape=(2, 2), strides=(1, 1)),
layer_values)
layer_values = add_layer(net,
snt.Linear(num_classes, **_inits),
layer_values,
flatten=True)
return layer_values
def __init__(self, name='MNIST_Discriminator', regularization=1.e-4):
super(MNISTDiscriminator, self).__init__(name=name)
reg = {
'w': l2_regularizer(scale=regularization),
'b': l2_regularizer(scale=regularization)
}
with self._enter_variable_scope():
self.conv1 = snt.Conv2D(name='conv2d_1',
output_channels=8,
kernel_shape=5,
stride=2,
regularizers=reg)
self.bn1 = snt.BatchNorm(name='batch_norm_1')
self.conv2 = snt.Conv2D(name='conv2d_2',
output_channels=16,
kernel_shape=5,
stride=1,
regularizers=reg)
self.bn2 = snt.BatchNorm(name='batch_norm_2')
self.conv3 = snt.Conv2D(name='conv2d_3',
output_channels=32,
kernel_shape=5,
stride=2,
regularizers=reg)
self.bn3 = snt.BatchNorm(name='batch_norm_3')
self.conv4 = snt.Conv2D(name='conv2d_4',
output_channels=64,
kernel_shape=5,
stride=1,
regularizers=reg)
self.bn4 = snt.BatchNorm(name='batch_norm_4')
self.conv5 = snt.Conv2D(name='conv2d_5',
output_channels=65,
kernel_shape=5,
stride=2,
regularizers=reg)
self.bn5 = snt.BatchNorm(name='batch_norm_5')
self.flatten = snt.BatchFlatten(name='flatten')
self.linear = snt.Linear(name='l', output_size=1, regularizers=reg)
def __init__(self, name='Discriminator', image_size=64, ndf=64, regularization=1.e-4):
super(Discriminator, self).__init__(name=name)
reg = {'w': l2_regularizer(scale=regularization)}
self.convs = []
self.batch_norms = []
csize, cndf = image_size / 2, ndf
with self._enter_variable_scope():
self.convs.append(snt.Conv2D(name='conv2d_1',
output_channels=ndf,
kernel_shape=4,
stride=2,
padding='VALID',
regularizers=reg,
use_bias=False))
self.batch_norms.append(snt.BatchNorm(name='batch_norm_1'))
n_layer = 2
while csize > 4:
self.convs.append(snt.Conv2D(name='conv2d_{}'.format(n_layer),
output_channels=cndf * 2,
kernel_shape=4,
stride=2,
padding='SAME',
regularizers=reg,
use_bias=False))
self.batch_norms.append(snt.BatchNorm(name='batch_norm_{}'.format(n_layer)))
cndf = cndf * 2
csize = csize // 2
n_layer += 1
self.convs.append(snt.Conv2D(name='conv2d_{}'.format(n_layer),
output_channels=1,
kernel_shape=4,
stride=1,
padding='SAME',
regularizers=reg,
use_bias=False))
def _build(self, h):
for i in range(self._num_residual_layers):
h_i = self._activation(h)
h_i = snt.Conv2D(output_channels=self._num_residual_hiddens,
kernel_shape=(self._filter_size,
self._filter_size),
stride=(1, 1),
initializers=self._initializers,
data_format=self._data_format,
name='res_nxn_%d' % i)(h_i)
h_i = self._activation(h_i)
h_i = snt.Conv2D(output_channels=self._num_hiddens,
kernel_shape=(1, 1),
stride=(1, 1),
initializers=self._initializers,
data_format=self._data_format,
name='res_1x1_%d' % i)(h_i)
h += h_i
return self._activation(h)
def testGetNormalizedVariableMapModule(self):
input_ = tf.placeholder(tf.float32, shape=[1, 10, 10, 3])
conv = snt.Conv2D(output_channels=3, kernel_shape=3)
conv(input_)
variable_map = snt.get_normalized_variable_map(conv)
self.assertEqual(len(variable_map), 2)
self.assertIn("w", variable_map)
self.assertIn("b", variable_map)
self.assertIs(variable_map["w"], conv.w)
self.assertIs(variable_map["b"], conv.b)
def __init__(self,
action_spec: specs.DiscreteArray,
name: Optional[Text] = None):
super().__init__(name=name)
# Spatial
self.conv1 = snt.Conv2D(16, 1, 1, data_format="NHWC", name="conv_1")
self.conv2 = snt.Conv2D(32, 3, 1, data_format="NHWC", name="conv_2")
self.conv3 = snt.Conv2D(64, 3, 1, data_format="NHWC", name="conv_3")
self.conv4 = snt.Conv2D(32, 3, 1, data_format="NHWC", name="conv_4")
self.flatten = snt.Flatten()
self.fc1 = snt.Linear(256, name="fc_1")
# Flat
self.flat = snt.nets.MLP([64, 64], name="mlp_1")
self.rnn = snt.DeepRNN([
snt.nets.MLP([50, 50], activate_final=True, name="mlp_2"),
snt.GRU(512, name="gru"),
networks.PolicyValueHead(action_spec.num_values)
])
def _build(self, x):
h = snt.Conv2D(output_channels=self._num_hiddens / 2,
kernel_shape=(4, 4),
stride=(2, 2),
name="enc_1")(x)
h = tf.nn.relu(h)
h = snt.Conv2D(output_channels=self._num_hiddens,
kernel_shape=(4, 4),
stride=(2, 2),
name="enc_2")(h)
h = tf.nn.relu(h)
h = snt.Conv2D(output_channels=self._num_hiddens,
kernel_shape=(3, 3),
stride=(1, 1),
name="enc_3")(h)
h = residual_stack(h, self._num_hiddens, self._num_residual_layers,
self._num_residual_hiddens)
return h
def test_get_parameter_overview_empty(self):
module = snt.Module()
snt.allow_empty_variables(module)
# No variables.
self.assertEqual(EMPTY_PARAMETER_OVERVIEW,
parameter_overview.get_parameter_overview(module))
module.conv = snt.Conv2D(output_channels=2, kernel_shape=3)
# Variables not yet created (happens in the first forward pass).
self.assertEqual(EMPTY_PARAMETER_OVERVIEW,
parameter_overview.get_parameter_overview(module))
def _build(self, x):
h = snt.Conv2D(output_channels=self._num_hiddens,
kernel_shape=(3, 3),
stride=(1, 1),
name="dec_1")(x)
h = self._dropout(h, training=self._is_training)
h = tf.layers.batch_normalization(
h,
training=self._is_training,
momentum=self._bn_momentum,
renorm=self._bn_renormalization,
renorm_momentum=self._bn_momentum,
renorm_clipping=self._renorm_clipping,
name="batch_norm_1")
h = residual_stack(h,
self._num_hiddens,
self._num_residual_layers,
self._num_residual_hiddens,
activation=self._activation,
training=self._is_training,
prob_drop=self._prob_drop,
momentum=self._bn_momentum,
renorm=self._bn_renormalization,
renorm_momentum=self._bn_momentum,
renorm_clipping=self._renorm_clipping)
h = snt.Conv2DTranspose(output_channels=int(self._num_hiddens / 2),
output_shape=None,
kernel_shape=(4, 4),
stride=(2, 2),
name="dec_2")(h)
h = self._dropout(h, training=self._is_training)
h = tf.layers.batch_normalization(
h,
training=self._is_training,
momentum=self._bn_momentum,
renorm=self._bn_renormalization,
renorm_momentum=self._bn_momentum,
renorm_clipping=self._renorm_clipping,
name="batch_norm_2")
h = self._activation(h)
x_recon = snt.Conv2DTranspose(output_channels=3,
output_shape=None,
kernel_shape=(4, 4),
stride=(2, 2),
name="dec_3")(h)
return x_recon
def _build(self, inputs, is_training=True, test_local_stats=False):
"""
Args:
inputs (type): node of input.
is_training (type): tells to batchnorm if to generate the update ops.
test_local_stats (type): used to test local stats in batch norm.
Returns:
logits
"""
# instantiate all the convolutional layers
self.layers = [
snt.Conv2D(name="conv_2d",
output_channels=self._conv_channels,
kernel_shape=self._conv_kernel_shape,
stride=self._conv_stride,
padding=self._padding,
use_bias=True,
**self._extra_params)
]
#(self, depth, name="resUnit", kernel_shape=[3,3], stride=1, activation=tf.nn.relu, **extra_params)
for i in range(self._num_resunits):
self.layers.append(
ResUnit(depth=self._resunit_channels[i],
name="resunit{}".format(i),
kernel_shape=self._resunit_kernel_shapes[i],
stride=self._resunit_strides[i],
activation=self._activation,
**self._extra_params))
net = self.layers[0](inputs)
net = tf.layers.max_pooling2d(net,
self._pooling_kernel_shape,
self._pooling_stride,
padding=self._padding,
data_format='channels_last',
name="max_pooling2d")
for i, resunit in enumerate(self.layers[1:]):
net = resunit(net)
net = tf.layers.average_pooling2d(net,
self._pooling_kernel_shape,
self._pooling_stride,
padding=self._padding,
data_format='channels_last',
name="avg_pooling2d")
net = snt.BatchFlatten(name="flatten")(net)
logits = snt.Linear(self._output_size)(net)
return logits
def test_get_parameter_overview_on_module(self):
module = snt.Module()
# Weights of a 2D convolution with 2 filters..
module.conv = snt.Conv2D(output_channels=2, kernel_shape=3, name="conv")
module.conv(tf.ones((2, 5, 5, 3))) # 3 * 3^2 * 2 = 56 parameters
for v in module.variables:
v.assign(tf.ones_like(v))
self.assertEqual(
SNT_CONV2D_PARAMETER_OVERVIEW,
parameter_overview.get_parameter_overview(module, include_stats=False))
self.assertEqual(SNT_CONV2D_PARAMETER_OVERVIEW_WITH_STATS,
parameter_overview.get_parameter_overview(module))
