Removed dependency on tensorflow.contrib for HParams.

PiperOrigin-RevId: 358219418
mbchang · Feb 18, 2021 · c374ee4 · c374ee4
1 parent 265d410
commit c374ee4
Show file tree

Hide file tree

Showing 6 changed files with 1,177 additions and 7 deletions.
diff --git a/graph_compression/compression_lib/compression_op.py b/graph_compression/compression_lib/compression_op.py
@@ -48,7 +48,7 @@
 import numpy as np
 import tensorflow.compat.v1 as tf
 from graph_compression.compression_lib import compression_op_utils as comp_op_utils
-from tensorflow.contrib import training as contrib_training
+from model_pruning.python import hparam as contrib_hparam
 
 
 class MatrixCompressorInferface(object):
@@ -131,7 +131,7 @@ def get_default_hparams():
     Returns:
       tf.HParams object initialized to default values.
     """
-    return contrib_training.HParams(
+    return contrib_hparam.HParams(
         name='model_compression',
         rank=100,
         num_rows=10,
@@ -333,7 +333,7 @@ def get_default_hparams():
       tf.HParams object initialized to default values.
 
     """
-    return contrib_training.HParams(
+    return contrib_hparam.HParams(
         name='model_compression',
         alpha_decrement_value=0.01,
         begin_compression_step=0,
@@ -889,7 +889,7 @@ def get_default_hparams():
       tf.HParams object initialized to default values.
 
     """
-    return contrib_training.HParams(
+    return contrib_hparam.HParams(
         name='input_compression',
         compression_frequency=10,
         use_tpu=False,

diff --git a/graph_compression/compression_lib/simhash_compression_op.py b/graph_compression/compression_lib/simhash_compression_op.py
@@ -27,8 +27,8 @@
 from graph_compression.compression_lib import compression_op
 from graph_compression.compression_lib import decompose_matrix
 from graph_compression.compression_lib import kmeans_quantize
+from model_pruning.python import hparam
 from model_pruning.python import pruning
-from tensorflow.contrib.training.python.training import hparam
 
 
 class SimhashMatrixCompressor(compression_op.LowRankDecompMatrixCompressor):

diff --git a/model_pruning/__init__.py b/model_pruning/__init__.py
@@ -20,6 +20,7 @@
 from __future__ import print_function
 
 # pylint: disable=unused-import
+from model_pruning.python.hparam import HParams
 from model_pruning.python.layers.rnn_cells import MaskedBasicLSTMCell
 from model_pruning.python.layers.rnn_cells import MaskedLSTMCell
 from model_pruning.python.pruning import apply_mask