update training driver

official/vision/beta/projects/basnet/common/registry_imports.py

-# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-"""All necessary imports for registration."""
-
-# pylint: disable=unused-import
-from official.vision import beta
-from official.vision.beta.projects.basnet.configs import basnet
-from official.vision.beta.projects.basnet.modeling import basnet_model
-from official.vision.beta.projects.basnet.modeling import refunet
-from official.vision.beta.projects.basnet.tasks import basnet

official/vision/beta/projects/basnet/train.py

-# Lint as: python3
-# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -12,61 +11,20 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-# ==============================================================================
-"""TensorFlow Model Garden Vision training driver."""
+
+# Lint as: python3
+"""TensorFlow Model Garden Vision training driver"""
 
 from absl import app
-from absl import flags
-import gin
 
-# pylint: disable=unused-import
-#from official.common import registry_imports
-from official.vision.beta.projects.basnet.common import registry_imports
-# pylint: enable=unused-import
-from official.common import distribute_utils
 from official.common import flags as tfm_flags
-from official.core import task_factory
-from official.core import train_lib
-from official.core import train_utils
-from official.modeling import performance
-
-FLAGS = flags.FLAGS
-
-
-def main(_):
-  gin.parse_config_files_and_bindings(FLAGS.gin_file, FLAGS.gin_params)
-  params = train_utils.parse_configuration(FLAGS)
-  model_dir = FLAGS.model_dir
-  if 'train' in FLAGS.mode:
-    # Pure eval modes do not output yaml files. Otherwise continuous eval job
-    # may race against the train job for writing the same file.
-    train_utils.serialize_config(params, model_dir)
-
-  # Sets mixed_precision policy. Using 'mixed_float16' or 'mixed_bfloat16'
-  # can have significant impact on model speeds by utilizing float16 in case of
-  # GPUs, and bfloat16 in the case of TPUs. loss_scale takes effect only when
-  # dtype is float16
-  if params.runtime.mixed_precision_dtype:
-    performance.set_mixed_precision_policy(params.runtime.mixed_precision_dtype,
-                                           params.runtime.loss_scale,
-                                           use_experimental_api=True)
-  distribution_strategy = distribute_utils.get_distribution_strategy(
-      distribution_strategy=params.runtime.distribution_strategy,
-      all_reduce_alg=params.runtime.all_reduce_alg,
-      num_gpus=params.runtime.num_gpus,
-      tpu_address=params.runtime.tpu)
-  with distribution_strategy.scope():
-    task = task_factory.get_task(params.task, logging_dir=model_dir)
-
-  train_lib.run_experiment(
-      distribution_strategy=distribution_strategy,
-      task=task,
-      mode=FLAGS.mode,
-      params=params,
-      model_dir=model_dir)
+from official.vision.beta import train
+from official.vision.beta.projects.basnet.configs import basnet
+from official.vision.beta.projects.basnet.modeling import basnet_model
+from official.vision.beta.projects.basnet.modeling import refunet
+from official.vision.beta.projects.basnet.tasks import basnet
 
-  train_utils.save_gin_config(FLAGS.mode, model_dir)
 
 if __name__ == '__main__':
   tfm_flags.define_flags()
-  app.run(main)
+  app.run(train.main)