Move Estimator-only export utils to R1

PiperOrigin-RevId: 264300408

official/r1/resnet/resnet_run_loop.py

@@ -33,7 +33,7 @@ import tensorflow as tf
 
 from official.r1.resnet import imagenet_preprocessing
 from official.r1.resnet import resnet_model
-from official.utils.export import export
+from official.r1.utils import export
 from official.utils.flags import core as flags_core
 from official.utils.logs import hooks_helper
 from official.utils.logs import logger

official/utils/export/export_test.py → official/r1/utils/export_test.py

@@ -20,7 +20,7 @@ from __future__ import print_function
 
 import tensorflow as tf  # pylint: disable=g-bad-import-order
 
-from official.utils.export import export
+from official.r1.utils import export
 
 
 class ExportUtilsTest(tf.test.TestCase):

official/transformer/transformer_main.py

@@ -32,6 +32,7 @@ from absl import flags
 import tensorflow as tf
 # pylint: enable=g-bad-import-order
 
+from official.r1.utils import export
 from official.transformer import compute_bleu
 from official.transformer import translate
 from official.transformer.model import model_params
@@ -41,7 +42,6 @@ from official.transformer.utils import metrics
 from official.transformer.utils import schedule
 from official.transformer.utils import tokenizer
 from official.utils.accelerator import tpu as tpu_util
-from official.utils.export import export
 from official.utils.flags import core as flags_core
 from official.utils.logs import hooks_helper
 from official.utils.logs import logger
@@ -56,7 +56,7 @@ PARAMS_MAP = {
 
 
 DEFAULT_TRAIN_EPOCHS = 10
-INF = int(1e9)
+INF = 1000000000  # 1e9
 BLEU_DIR = "bleu"
 
 # Dictionary containing tensors that are logged by the logging hooks. Each item