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