def testLowRankDecompMatrixCompressorInterface(self): spec = compression_op.LowRankDecompMatrixCompressor.get_default_hparams() compressor = compression_op.LowRankDecompMatrixCompressor(spec) b_matrix = np.random.normal(0, 1, [10, 5]) c_matrix = np.random.normal(0, 1, [5, 10]) a_matrix = np.matmul(b_matrix, c_matrix) [b_matrix_out, c_matrix_out] = compressor.static_matrix_compressor(a_matrix) a_matrix_recovered = np.matmul(b_matrix_out, c_matrix_out) self.assertLess(np.linalg.norm(a_matrix - a_matrix_recovered), 0.01)
def test_get_apply_matmul(self):
with tf.Graph().as_default():
with self.cached_session():
hparams = ("name=block_compression,"
"compression_option=10,"
"begin_compression_step=1000,"
"end_compression_step=120000,"
"compression_frequency=100,"
"block_method=mask,"
"block_compression_factor=2,")
compression_op_spec = (compression_op.BlockCompressionOp.
get_default_hparams().parse(hparams))
compressor_spec = (compression_op.LowRankDecompMatrixCompressor
.get_default_hparams())
matrix_compressor = compression_op.LowRankDecompMatrixCompressor(
spec=compressor_spec)
global_step = tf.compat.v1.get_variable("global_step",
initializer=100)
apply_comp = compression_op.ApplyCompression(
scope="default_scope",
compression_spec=compression_op_spec,
compressor=matrix_compressor,
global_step=global_step)
# outer product - creates an 12x8 matrix
a_matrix_init = np.outer(
np.array(
[1., 2., 3., 7., 8., 9., 1., 2., 5., -2., -7., -1.]),
np.array([4., 5., 6., 3., 1., 8., 3., 2.]))
a_matrix = tf.compat.v1.get_variable(
"a_matrix",
initializer=a_matrix_init.astype(np.float32),
dtype=tf.float32)
_ = apply_comp.apply_compression(a_matrix, scope="compressor")
# input is 1x12 vector
left_operand_init = np.expand_dims(np.array(
[1., 2., 3., 4., 1., 2., 3., 4., 1., 2., 3., 4.]),
axis=0)
left_operand = tf.compat.v1.get_variable(
"left_operand",
initializer=left_operand_init.astype(np.float32),
dtype=tf.float32)
c = apply_comp._compression_ops[-1]
tf.compat.v1.global_variables_initializer().run()
compressed_matmul = c.get_apply_matmul(left_operand)
# check c, c_mask matrices have the right shapes
self.assertSequenceEqual(list(c.c_matrix_tfvar.eval().shape),
[12, 8])
self.assertSequenceEqual(list(c.c_mask_tfvar.eval().shape),
[12, 8])
# check we get the correct number of nonzero entries in the mask
self.assertEqual(np.count_nonzero(c.c_mask_tfvar.eval()), 48)
# check that we get the expected output shape
self.assertSequenceEqual(list(compressed_matmul.eval().shape),
[1, 8])
def test_get_apply_matmul(self):
with tf.Graph().as_default():
with self.cached_session():
hparams = ("name=input_output_compression,"
"compression_option=9,"
"begin_compression_step=1000,"
"end_compression_step=120000,"
"compression_frequency=100,"
"compress_input=True,"
"compress_output=True,"
"input_compression_factor=2,"
"input_block_size=4,"
"output_compression_factor=2,"
"output_block_size=4,")
compression_op_spec = (
compression_op.InputOutputCompressionOp.get_default_hparams().parse(
hparams))
compressor_spec = (
compression_op.LowRankDecompMatrixCompressor.get_default_hparams())
matrix_compressor = compression_op.LowRankDecompMatrixCompressor(
spec=compressor_spec)
global_step = tf.compat.v1.get_variable("global_step", initializer=100)
apply_comp = compression_op.ApplyCompression(
scope="default_scope",
compression_spec=compression_op_spec,
compressor=matrix_compressor,
global_step=global_step)
# outer product - creates an 12x8 matrix
a_matrix_init = np.outer(
np.array([1., 2., 3., 7., 8., 9., 1., 2., 5., -2., -7., -1.]),
np.array([4., 5., 6., 3., 1., 8., 3., 2.]))
a_matrix = tf.compat.v1.get_variable(
"a_matrix",
initializer=a_matrix_init.astype(np.float32),
dtype=tf.float32)
_ = apply_comp.apply_compression(
a_matrix, scope="compressor")
# input is 1x12 vector
left_operand_init = np.array(
[1., 2., 3., 4., 1., 2., 3., 4., 1., 2., 3., 4.])
left_operand = tf.compat.v1.get_variable(
"left_operand",
initializer=left_operand_init.astype(np.float32),
dtype=tf.float32)
c = apply_comp._compression_ops[-1]
tf.compat.v1.global_variables_initializer().run()
compressed_matmul = c.get_apply_matmul(left_operand)
# check b, c and d matrices have the right shapes
self.assertSequenceEqual(list(c.b_matrix_tfvar.eval().shape), [4, 2])
self.assertSequenceEqual(list(c.c_matrix_tfvar.eval().shape), [6, 4])
self.assertSequenceEqual(list(c.d_matrix_tfvar.eval().shape), [2, 4])
# check that we get the expected output shape
self.assertSequenceEqual(list(compressed_matmul.eval().shape), [8,])
def get_apply_compression(compression_op_spec, global_step):
"""Returns apply_compression operation matching compression_option input."""
compressor_spec = comp_op.LowRankDecompMatrixCompressor.get_default_hparams(
)
compressor_spec.set_hparam('rank', compression_op_spec.rank)
compressor_spec.set_hparam('block_size', compression_op_spec.block_size)
logging.info('Compressor spec %s', compressor_spec.to_json())
logging.info('Compression operator spec %s', compression_op_spec.to_json())
if compression_op_spec.compression_option not in _COMPRESSION_OPTIONS:
logging.info(
'Compression_option %s not in expected options: %s. '
'Will use low_rank decomp by default.',
str(compression_op_spec.compression_option),
','.join([str(opt) for opt in _COMPRESSION_OPTIONS]))
compression_op_spec.compression_option = 1
apply_compression = None
if compression_op_spec.compression_option == 1:
compressor = comp_op.LowRankDecompMatrixCompressor(
spec=compressor_spec)
apply_compression = comp_op.ApplyCompression(
scope='default_scope',
compression_spec=compression_op_spec,
compressor=compressor,
global_step=global_step)
elif compression_op_spec.compression_option == 2:
compressor_spec.set_hparam('is_b_matrix_trainable', False)
compressor = simhash_comp_op.SimhashMatrixCompressor(
spec=compressor_spec)
apply_compression = simhash_comp_op.SimhashApplyCompression(
scope='default_scope',
compression_spec=compression_op_spec,
compressor=compressor,
global_step=global_step)
elif compression_op_spec.compression_option == 4:
compressor_spec.set_hparam('is_b_matrix_trainable', False)
compressor = simhash_comp_op.KmeansMatrixCompressor(
spec=compressor_spec)
apply_compression = simhash_comp_op.SimhashApplyCompression(
scope='default_scope',
compression_spec=compression_op_spec,
compressor=compressor,
global_step=global_step)
elif compression_op_spec.compression_option == 8:
compressor_spec.set_hparam('is_b_matrix_trainable', False)
compressor = simhash_comp_op.KmeansMatrixCompressor(
spec=compressor_spec)
apply_compression = simhash_comp_op.SimhashApplyCompression(
scope='default_scope',
compression_spec=compression_op_spec,
compressor=compressor,
global_step=global_step)
return apply_compression
def test_get_apply_matmul(self):
with tf.Graph().as_default():
with self.cached_session():
hparams = ("name=mixed_block_compression,"
"begin_compression_step=1000,"
"end_compression_step=120000,"
"compression_frequency=100,"
"compression_factor=4,"
"num_bases=2,")
compression_op_spec = (compression_op.MixedBlockCompressionOp.
get_default_hparams().parse(hparams))
compressor_spec = (compression_op.LowRankDecompMatrixCompressor
.get_default_hparams())
matrix_compressor = compression_op.LowRankDecompMatrixCompressor(
spec=compressor_spec)
global_step = tf.compat.v1.get_variable("global_step",
initializer=100)
apply_comp = compression_wrapper.ApplyCompression(
scope="default_scope",
compression_spec=compression_op_spec,
compressor=matrix_compressor,
global_step=global_step)
# outer product - creates an 12x8 matrix
a_matrix_init = np.outer(
np.array(
[1., 2., 3., 7., 8., 9., 1., 2., 5., -2., -7., -1.]),
np.array([4., 5., 6., 3., 1., 8., 3., 2.]))
a_matrix = tf.compat.v1.get_variable(
"a_matrix",
initializer=a_matrix_init.astype(np.float32),
dtype=tf.float32)
_ = apply_comp.apply_compression(a_matrix, scope="compressor")
# input is 1x12 vector
left_operand_init = np.expand_dims(np.array(
[1., 2., 3., 4., 1., 2., 3., 4., 1., 2., 3., 4.]),
axis=0)
left_operand = tf.compat.v1.get_variable(
"left_operand",
initializer=left_operand_init.astype(np.float32),
dtype=tf.float32)
c = apply_comp._compression_ops[-1]
tf.compat.v1.global_variables_initializer().run()
compressed_matmul = c.get_apply_matmul(left_operand)
# check block_matrices and linear_mixer tensors have the right shapes
self.assertSequenceEqual(list(c.block_matrices.eval().shape),
[3, 2, 4])
self.assertSequenceEqual(list(c.linear_mixer.eval().shape),
[4, 4, 2])
# check that we get the expected output shape
self.assertSequenceEqual(list(compressed_matmul.eval().shape),
[1, 8])
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
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.
input_dim = 28 * 28
num_hidden_nodes = 50
num_classes = 10
lowrank_compressor = compression.LowRankDecompMatrixCompressor(
compression.LowRankDecompMatrixCompressor.get_default_hparams())
compressed_model = CompressedModel(input_dim, num_hidden_nodes,
num_classes, lowrank_compressor)
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 testCompressionOpInterface(self):
with tf.Graph().as_default():
with self.cached_session() as sess:
compression_hparams = ("name=cifar10_compression,"
"begin_compression_step=1000,"
"end_compression_step=120000,"
"compression_frequency=10,"
"compression_option=1,"
"update_option=0")
global_step = tf.compat.v1.get_variable("global_step", initializer=30)
c = compression_op.CompressionOp(
spec=compression_op.CompressionOp.get_default_hparams().parse(
compression_hparams),
global_step=global_step)
# Need to add initial value for a_matrix so that we would know what
# to expect back.
a_matrix_init = np.array([[1.0, 1.0, 1.0], [1.0, 0, 0], [1.0, 0, 0]])
a_matrix = tf.compat.v1.get_variable(
"a_matrix",
initializer=a_matrix_init.astype(np.float32),
dtype=tf.float32)
matrix_compressor = compression_op.LowRankDecompMatrixCompressor(
spec=compression_op.LowRankDecompMatrixCompressor
.get_default_hparams().parse("num_rows=3,num_cols=3,rank=200"))
[a_matrix_compressed, a_matrix_update_op] = c.get_apply_compression_op(
a_matrix, matrix_compressor, scope="my_scope")
tf.compat.v1.global_variables_initializer().run()
self.assertAllEqual(
np.all(np.abs(np.linalg.norm(c.a_matrix_tfvar.eval())) < 0.00001),
False)
self.assertAllEqual(
np.all(np.abs(np.linalg.norm(c.b_matrix_tfvar.eval())) < 0.00001),
True)
self.assertAllEqual(
np.all(np.abs(np.linalg.norm(c.c_matrix_tfvar.eval())) < 0.00001),
True)
tf.compat.v1.assign(global_step, 1001).eval()
sess.run(a_matrix_update_op)
a_matrix_compressed.eval()
self.assertEqual(c._global_step.eval(), 1001)
self.assertAlmostEqual(c.alpha.eval(), 0.99)
self.assertEqual(c._last_alpha_update_step.eval(), 1001)
self.assertAllEqual(
np.array([
np.linalg.norm(c.a_matrix_tfvar.eval()),
np.linalg.norm(c.b_matrix_tfvar.eval()),
np.linalg.norm(c.c_matrix_tfvar.eval())
]) > 0, [True, True, True])
self.assertAllEqual(
np.all(np.abs(np.linalg.norm(c.b_matrix_tfvar.eval())) < 0.00001),
False)
self.assertAllEqual(
np.all(np.abs(np.linalg.norm(c.c_matrix_tfvar.eval())) < 0.00001),
False)
[b_matrix,
c_matrix] = matrix_compressor.static_matrix_compressor(a_matrix_init)
# since the matrices may match up to signs, we take absolute values.
self.assertAllEqual(
np.linalg.norm(np.abs(b_matrix) - np.abs(c.b_matrix_tfvar.eval())) <
0.00001, True)
self.assertAllEqual(
np.linalg.norm(np.abs(c_matrix) - np.abs(c.c_matrix_tfvar.eval())) <
0.00001, True)
self.assertAllEqual(
np.all(np.abs(np.linalg.norm(c.b_matrix_tfvar.eval())) < 0.00001),
False)
self.assertAllEqual(
np.all(np.abs(np.linalg.norm(c.c_matrix_tfvar.eval())) < 0.00001),
False)
print("before 1002 step, c.alpha is ", c.alpha.eval())
tf.compat.v1.assign(global_step, 1001).eval()
sess.run(a_matrix_update_op)
a_matrix_compressed.eval()
print("after 1002 step, c.alpha is ", c.alpha.eval())
self.assertEqual(c._global_step.eval(), 1001)
self.assertAlmostEqual(c.alpha.eval(), 0.99)
self.assertEqual(c._last_alpha_update_step.eval(), 1001)
self.assertAllEqual(
np.all([
np.linalg.norm(c.a_matrix_tfvar.eval()),
np.linalg.norm(c.b_matrix_tfvar.eval()),
np.linalg.norm(c.c_matrix_tfvar.eval())
]) > 0, True)
print("before 2000 step, alpha is ", c.alpha.eval())
tf.compat.v1.assign(global_step, 2000).eval()
a_matrix_update_op.eval()
a_matrix_compressed.eval()
print("after 2000 step, alpha is ", c.alpha.eval())
self.assertEqual(c._global_step.eval(), 2000)
self.assertAlmostEqual(c.alpha.eval(), 0.98)
self.assertEqual(c._last_alpha_update_step.eval(), 2000)
self.assertAllEqual(
np.array([
np.linalg.norm(c.a_matrix_tfvar.eval()),
np.linalg.norm(c.b_matrix_tfvar.eval()),
np.linalg.norm(c.c_matrix_tfvar.eval())
]) > 0, [True, True, True])
def testApplyCompression(self):
with tf.Graph().as_default():
with self.cached_session():
compression_hparams = ("name=cifar10_compression,"
"begin_compression_step=1000,"
"end_compression_step=120000,"
"compression_frequency=100,"
"compression_option=1")
compression_op_spec = (
compression_op.CompressionOp.get_default_hparams().parse(
compression_hparams))
compressor_spec = (
compression_op.LowRankDecompMatrixCompressor.get_default_hparams()
.parse("num_rows=5,num_cols=5,rank=200"))
matrix_compressor = compression_op.LowRankDecompMatrixCompressor(
spec=compressor_spec)
global_step = tf.compat.v1.get_variable("global_step", initializer=30)
apply_comp = compression_op.ApplyCompression(
scope="default_scope",
compression_spec=compression_op_spec,
compressor=matrix_compressor,
global_step=global_step)
# Need to add initial value for a_matrix so that we would know what
# to expect back.
a_matrix_init = np.outer(np.array([1., 2., 3.]), np.array([4., 5., 6.]))
a_matrix = tf.compat.v1.get_variable(
"a_matrix",
initializer=a_matrix_init.astype(np.float32),
dtype=tf.float32)
a_matrix_compressed = apply_comp.apply_compression(
a_matrix, scope="first_compressor")
c = apply_comp._compression_ops[0]
a_matrix2 = tf.compat.v1.get_variable(
"a_matrix2",
initializer=a_matrix_init.astype(np.float32),
dtype=tf.float32)
_ = apply_comp.apply_compression(a_matrix2, scope="second_compressor")
c2 = apply_comp._compression_ops[1]
_ = apply_comp.all_update_op()
tf.compat.v1.global_variables_initializer().run()
_ = a_matrix_compressed.eval()
self.assertEqual(c._global_step.eval(), 30)
self.assertEqual(c.alpha.eval(), 1.0)
self.assertEqual(c2.alpha.eval(), 1.0)
self.assertEqual(c._last_alpha_update_step.eval(), -1)
self.assertAllEqual(
np.array([
np.linalg.norm(c.a_matrix_tfvar.eval()),
np.linalg.norm(c.b_matrix_tfvar.eval()),
np.linalg.norm(c.c_matrix_tfvar.eval())
]) > 0, [True, False, False])
self.assertAllEqual(
np.all(np.abs(np.linalg.norm(c.a_matrix_tfvar.eval())) < 0.00001),
False)
self.assertAllEqual(
np.all(np.abs(np.linalg.norm(c.b_matrix_tfvar.eval())) < 0.00001),
True)
self.assertAllEqual(
np.all(np.abs(np.linalg.norm(c.c_matrix_tfvar.eval())) < 0.00001),
True)
tf.compat.v1.assign(global_step, 1001).eval()
# apply_comp_update_op.run()
apply_comp._all_update_op.run()
_ = a_matrix_compressed.eval()
self.assertEqual(c._global_step.eval(), 1001)
self.assertAlmostEqual(c.alpha.eval(), 0.99)
self.assertEqual(c._last_alpha_update_step.eval(), 1001)
self.assertAllEqual(
np.array([
np.linalg.norm(c.a_matrix_tfvar.eval()),
np.linalg.norm(c.b_matrix_tfvar.eval()),
np.linalg.norm(c.c_matrix_tfvar.eval())
]) > 0, [True, True, True])
self.assertAllEqual(
np.all(np.abs(np.linalg.norm(c.b_matrix_tfvar.eval())) < 0.00001),
False)
self.assertAllEqual(
np.all(np.abs(np.linalg.norm(c.c_matrix_tfvar.eval())) < 0.00001),
False)
[b_matrix,
c_matrix] = matrix_compressor.static_matrix_compressor(a_matrix_init)
self.assertAllEqual(
np.linalg.norm(np.abs(b_matrix) - np.abs(c.b_matrix_tfvar.eval())) <
0.00001, True)
self.assertAllEqual(
np.linalg.norm(np.abs(c_matrix) - np.abs(c.c_matrix_tfvar.eval())) <
0.00001, True)
self.assertAllEqual(
np.all(np.abs(np.linalg.norm(c.b_matrix_tfvar.eval())) < 0.00001),
False)
self.assertAllEqual(
np.all(np.abs(np.linalg.norm(c.c_matrix_tfvar.eval())) < 0.00001),
False)
tf.compat.v1.assign(global_step, 1001).eval()
apply_comp._all_update_op.run()
_ = a_matrix_compressed.eval()
self.assertEqual(c._global_step.eval(), 1001)
self.assertAlmostEqual(c.alpha.eval(), 0.99)
self.assertEqual(c._last_alpha_update_step.eval(), 1001)
self.assertAllEqual(
np.array([
np.linalg.norm(c.a_matrix_tfvar.eval()),
np.linalg.norm(c.b_matrix_tfvar.eval()),
np.linalg.norm(c.c_matrix_tfvar.eval())
]) > 0, [True, True, True])
tf.compat.v1.assign(global_step, 2001).eval()
apply_comp._all_update_op.run()
_ = a_matrix_compressed.eval()
self.assertEqual(c._global_step.eval(), 2001)
self.assertAlmostEqual(c.alpha.eval(), 0.98)
self.assertAlmostEqual(c2.alpha.eval(), 0.98)
self.assertEqual(c._last_alpha_update_step.eval(), 2001)
self.assertAllEqual(
np.array([
np.linalg.norm(c.a_matrix_tfvar.eval()),
np.linalg.norm(c.b_matrix_tfvar.eval()),
np.linalg.norm(c.c_matrix_tfvar.eval())
]) > 0, [True, True, True])
def get_apply_compression(compression_op_spec, global_step):
"""Returns apply_compression operation matching compression_option input."""
compressor_spec = comp_op.LowRankDecompMatrixCompressor.get_default_hparams()
if compression_op_spec.__contains__('rank'):
compressor_spec.set_hparam('rank', compression_op_spec.rank)
if compression_op_spec.__contains__('block_size'):
compressor_spec.set_hparam('block_size', compression_op_spec.block_size)
logging.info('Compressor spec %s', compressor_spec.to_json())
logging.info('Compression operator spec %s', compression_op_spec.to_json())
if compression_op_spec.compression_option not in list(CompressionOptions):
# if unknown compression_option is given, default to low rank compression.
logging.info(
'Compression_option %s not in expected options: %s. '
'Will use low_rank decomp by default.',
str(compression_op_spec.compression_option),
','.join([str(opt) for opt in CompressionOptions]))
compression_op_spec.compression_option = CompressionOptions.LOWRANK_MATRIX_COMPRESSION
apply_compression = None
if compression_op_spec.compression_option == CompressionOptions.LOWRANK_MATRIX_COMPRESSION:
compressor = comp_op.LowRankDecompMatrixCompressor(spec=compressor_spec)
apply_compression = ApplyCompression(
scope='default_scope',
compression_spec=compression_op_spec,
compressor=compressor,
global_step=global_step)
elif compression_op_spec.compression_option == CompressionOptions.SIMHASH_MATRIX_COMPRESSION:
compressor_spec.set_hparam('is_b_matrix_trainable', False)
compressor = simhash_comp_op.SimhashMatrixCompressor(spec=compressor_spec)
apply_compression = ApplyCompression(
scope='default_scope',
compression_spec=compression_op_spec,
compressor=compressor,
global_step=global_step)
elif compression_op_spec.compression_option == CompressionOptions.KMEANS_MATRIX_COMPRESSION:
compressor_spec.set_hparam('is_b_matrix_trainable', True)
compressor = simhash_comp_op.KmeansMatrixCompressor(spec=compressor_spec)
apply_compression = ApplyCompression(
scope='default_scope',
compression_spec=compression_op_spec,
compressor=compressor,
global_step=global_step)
elif compression_op_spec.compression_option == CompressionOptions.KMEANS_AND_PRUNING_MATRIX_COMPRESSION:
compressor_spec.set_hparam('is_b_matrix_trainable', True)
compressor = simhash_comp_op.KmeansMatrixCompressor(spec=compressor_spec)
apply_compression = ApplyCompression(
scope='default_scope',
compression_spec=compression_op_spec,
compressor=compressor,
global_step=global_step)
elif compression_op_spec.compression_option == CompressionOptions.INPUTOUTPUT_COMPRESSION:
compressor_spec.set_hparam('is_b_matrix_trainable', True)
compressor_spec.set_hparam('is_c_matrix_trainable', True)
compressor_spec.set_hparam('is_d_matrix_trainable', True)
compressor = comp_op.LowRankDecompMatrixCompressor(spec=compressor_spec)
apply_compression = ApplyCompression(
scope='default_scope',
compression_spec=compression_op_spec,
compressor=compressor,
global_step=global_step)
elif compression_op_spec.compression_option == CompressionOptions.BLOCK_COMPRESSION:
compressor_spec.set_hparam('is_c_matrix_trainable', True)
compressor = comp_op.LowRankDecompMatrixCompressor(spec=compressor_spec)
apply_compression = ApplyCompression(
scope='default_scope',
compression_spec=compression_op_spec,
compressor=compressor,
global_step=global_step)
elif compression_op_spec.compression_option == CompressionOptions.MIXED_BLOCK_COMPRESSION:
compressor_spec.set_hparam('is_c_matrix_trainable', True)
compressor = comp_op.LowRankDecompMatrixCompressor(spec=compressor_spec)
apply_compression = ApplyCompression(
scope='default_scope',
compression_spec=compression_op_spec,
compressor=compressor,
global_step=global_step)
elif compression_op_spec.compression_option == CompressionOptions.DL_MATRIX_COMPRESSION:
compressor = dl_comp_op.DLMatrixCompressor(spec=compressor_spec)
apply_compression = ApplyCompression(
scope='default_scope',
compression_spec=compression_op_spec,
compressor=compressor,
global_step=global_step)
return apply_compression