def compressed_lenet5(input_shape, num_classes, compression_obj):
"""Builds Compressed version of LeNet5."""
inputs = tf.keras.layers.Input(shape=input_shape)
conv1 = compression_layers.CompressedConv2D(
6,
kernel_size=5,
padding='SAME',
activation='relu',
compression_obj=compression_obj)(inputs)
pool1 = tf.keras.layers.MaxPooling2D(pool_size=[2, 2],
strides=[2, 2],
padding='SAME')(conv1)
conv2 = compression_layers.CompressedConv2D(
16,
kernel_size=5,
padding='SAME',
activation='relu',
compression_obj=compression_obj)(pool1)
pool2 = tf.keras.layers.MaxPooling2D(pool_size=[2, 2],
strides=[2, 2],
padding='SAME')(conv2)
conv3 = compression_layers.CompressedConv2D(
120,
kernel_size=5,
padding='SAME',
activation=tf.nn.relu,
compression_obj=compression_obj)(pool2)
flatten = tf.keras.layers.Flatten()(conv3)
dense1 = compression_layers.CompressedDense(
84, activation=tf.nn.relu, compression_obj=compression_obj)(flatten)
logits = tf.keras.layers.Dense(num_classes)(dense1)
outputs = tf.keras.layers.Softmax()(logits)
return tf.keras.Model(inputs=inputs, outputs=outputs)
def testCompressedConv2DLayer(self):
hparams = ("name=mnist_compression,"
"compress_input=True,"
"input_block_size=16,"
"input_compression_factor=2,"
"compression_option=9")
compression_hparams = compression_op.InputOutputCompressionOp.get_default_hparams(
).parse(hparams)
compression_obj = compression_wrapper.get_apply_compression(
compression_hparams, global_step=0)
val = np.random.random((10, 4, 10, 10))
x = tf.Variable(val, dtype=tf.float32)
y_compressed = compression_layers.CompressedConv2D(
20,
3,
padding="valid",
data_format="channels_last",
compression_obj=compression_obj)(x)
y = tf.keras.layers.Conv2D(20,
3,
padding="valid",
data_format="channels_last")(x)
self.assertAllEqual(y.shape.as_list(), y_compressed.shape.as_list())