def testWrapper_CreatesProperCompressorOption1(self, low_rank_mock):
hparams = self._create_compression_op_spec(1)
mock_compressor = MatrixCompressorInterfaceMock(
self._default_compressor_spec(hparams))
low_rank_mock.side_effect = [mock_compressor]
with mock.patch.object(comp_op, 'ApplyCompression') as apply_mock:
compression_wrapper.get_apply_compression(hparams, _GLOBAL_STEP)
apply_mock.assert_called_with(scope='default_scope',
compression_spec=hparams,
compressor=mock_compressor,
global_step=_GLOBAL_STEP)
def testWrapper_CreatesProperCompressorOption2(self, sim_hash_mock):
hparams = self._create_compression_op_spec(
compression_wrapper.CompressionOptions.SIMHASH_MATRIX_COMPRESSION,
compression_wrapper.UpdateOptions.PYTHON_UPDATE)
mock_compressor = MatrixCompressorInterfaceMock(
self._default_compressor_spec(hparams))
sim_hash_mock.side_effect = [mock_compressor]
with mock.patch.object(compression_wrapper,
'ApplyCompression') as apply_mock:
compression_wrapper.get_apply_compression(hparams, _GLOBAL_STEP)
apply_mock.assert_called_with(scope='default_scope',
compression_spec=hparams,
compressor=mock_compressor,
global_step=_GLOBAL_STEP)
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())
def testWrapper_CreatesProperCompressorOption1(self):
hparams = self._create_compression_op_spec(
compression_wrapper.CompressionOptions.LOWRANK_MATRIX_COMPRESSION,
compression_wrapper.UpdateOptions.PYTHON_UPDATE)
mock_compressor = MatrixCompressorInterfaceMock(
self._default_compressor_spec(hparams))
self.enter_context(
mock.patch.object(comp_op,
'LowRankDecompMatrixCompressor',
side_effect=[mock_compressor]))
with mock.patch.object(compression_wrapper,
'ApplyCompression') as apply_mock:
compression_wrapper.get_apply_compression(hparams, _GLOBAL_STEP)
apply_mock.assert_called_with(scope='default_scope',
compression_spec=hparams,
compressor=mock_compressor,
global_step=_GLOBAL_STEP)
def main(argv):
del argv # unused
tf.enable_v2_behavior()
# Load MNIST data.
mnist = tf.keras.datasets.mnist
(x_train, y_train), (_, _) = mnist.load_data()
x_train = x_train / 255.0
if not FLAGS.use_lennet:
x_train = x_train.reshape(60000, 784).astype('float32')
train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
train_dataset = train_dataset.shuffle(buffer_size=1024).batch(64)
# Define model.
num_hidden_nodes = 64
num_classes = 10
hparams = ('name=mnist_compression,'
'prune_option=compression,'
'input_block_size=20,'
'rank=2,'
'compression_option=9')
compression_hparams = compression.InputCompressionOp.get_default_hparams(
).parse(hparams)
compression_obj = compression_wrapper.get_apply_compression(
compression_hparams, global_step=0)
if not FLAGS.use_lennet:
compressed_model = CompressedModelV2(num_hidden_nodes, num_classes,
compression_obj)
else:
compressed_model = compressed_lenet5([28, 28, 1], num_classes,
compression_obj)
optimizer = tf.keras.optimizers.SGD(learning_rate=1e-3)
loss = tf.keras.losses.SparseCategoricalCrossentropy()
epochs = 10
step_number = 0
for epoch in range(epochs):
for x, y in train_dataset:
with tf.GradientTape() as t:
loss_value = loss(y, compressed_model(x))
grads = t.gradient(loss_value,
compressed_model.trainable_variables)
optimizer.apply_gradients(
zip(grads, compressed_model.trainable_variables))
# compressed_model.run_alpha_update(step_number)
step_number += 1
print('Training loss at epoch {} is {}.'.format(epoch, loss_value))
def create_compressor(hparams='', global_step=None):
"""Creates an ApplyCompression object to use during the inference step."""
# Parse compression hyperparameters
compression_hparams = compression.CompressionOp.get_default_hparams().parse(
hparams)
# Create a compression object using the compression hyperparameters
compression_obj = compression_wrapper.get_apply_compression(
compression_hparams, global_step=global_step)
return compression_obj
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())
def get_matrix_compression_object(hparams,
global_step=None,
sparsity=None):
"""Returns a pruning/compression object.
Args:
hparams: Pruning spec as defined in pruing.py;
global_step: A tensorflow variable that is used for scheduling
pruning/compression;
sparsity: A tensorflow scalar variable storing the sparsity.
Returns:
A Pruning or compression_lib.compression_op.ApplyCompression object.
"""
if not global_step:
global_step = tf.cast(tf.train.get_global_step(), tf.int32)
if hparams.prune_option in [
'weight', 'first_order_gradient', 'second_order_gradient']:
return pruning.Pruning(hparams, global_step, sparsity)
else:
return compression_wrapper.get_apply_compression(hparams,
global_step=global_step)
