def __init__(self, num_hidden_nodes, num_classes, compression_obj,
compression_flag=True):
"""Initializer.
Args:
num_hidden_nodes: int
num_classes: int
compression_obj: a matrix compression object obtained by calling
compression_wrapper
compression_flag: if True compressed model will be used
"""
super().__init__()
self.num_hidden_nodes = num_hidden_nodes
self.num_classes = num_classes
self.compression_obj = compression_obj
compression_flag = True
if compression_flag:
self.layer_1 = compression_layers.CompressedDense(
num_hidden_nodes, compression_obj=compression_obj, activation='relu')
self.layer_2 = compression_layers.CompressedDense(
num_classes, compression_obj=compression_obj)
else:
self.layer_1 = tf.keras.layers.Dense(num_hidden_nodes, activation='relu')
self.layer_2 = tf.keras.layers.Dense(num_classes)
self.softmax = tf.keras.layers.Softmax()
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 testCompressedDenseLayer(self):
hparams = ("name=mnist_compression,"
"compress_input=True,"
"input_block_size=16,"
"input_compression_factor=4,")
compression_hparams = compression_op.InputOutputCompressionOp.get_default_hparams(
).parse(hparams)
# Create a compression object using the compression hyperparameters
compression_obj = compression_wrapper.get_apply_compression(
compression_hparams, global_step=0)
val = np.random.random((10, 48))
x = tf.Variable(val, dtype=tf.float32)
y_compressed = compression_layers.CompressedDense(
20, compression_obj=compression_obj)(x)
y = tf.keras.layers.Dense(20)(x)
self.assertAllEqual(y.shape.as_list(), y_compressed.shape.as_list())
