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())