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 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
def testSimhashApplyCompressionKmeans(self): with self.cached_session() as session: hparams = ("name=cifar10_compression," "begin_compression_step=1000," "end_compression_step=2001," "compression_frequency=100," "compression_option=4," "update_option=1") spec = simhash.SimhashCompressionOp.get_default_hparams().parse( hparams) matrix_compressor = simhash.KmeansMatrixCompressor( spec=compression_op.LowRankDecompMatrixCompressor. get_default_hparams().parse( "rank=2,block_size=3,pruning_fraction=0.0")) global_step = tf.get_variable("global_step", initializer=30) apply_comp = simhash.SimhashApplyCompression( scope="default_scope", compression_spec=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.])) jitter = np.tile([0, 1e-1, 2e-2], (3, 1)) a_matrix_init += jitter a_matrix = tf.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.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] tf.global_variables_initializer().run() _ = a_matrix_compressed.eval() # Compression won't start until step 1000 + some random_shift amount. # Here we make sure output is as expected at step 30. 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()), ]) > 0, [True, False]) self.assertAllEqual( np.array([ np.linalg.norm(c.a_matrix_tfvar.eval()), np.linalg.norm(c.b_matrix_tfvar.eval()), ]) < 0.00001, [False, True]) # At this point compression should have already started being applied; # verify at step 2000 all is as expected. tf.assign(global_step, 2000).eval() apply_comp.run_update_step(session=session) self.assertEqual(c._global_step.eval(), 2000) self.assertAlmostEqual(c.alpha.eval(), 0.99) 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()), ]) > 0, [True, True]) self.assertFalse( np.all( np.abs(np.linalg.norm(c.b_matrix_tfvar.eval())) < 0.00001)) [b_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)