From 9fd6313f37c3bc124933c582a99f1ed339c9a2a1 Mon Sep 17 00:00:00 2001 From: Scott Zhu Date: Tue, 6 Oct 2020 14:40:50 -0700 Subject: [PATCH] Export strategy combinations that used by keras to tf.__internal__. PiperOrigin-RevId: 335724503 Change-Id: I15e5adbba764d91a1d81ba60724719d678697e6f --- tensorflow/python/BUILD | 1 - .../distribute/strategy_combinations.py | 63 +---------------- tensorflow/python/modules_with_exports.py | 1 - ...__internal__.distribute.combinations.pbtxt | 68 ------------------- 4 files changed, 1 insertion(+), 132 deletions(-) diff --git a/tensorflow/python/BUILD b/tensorflow/python/BUILD index 9e2e481df27..7825f7e5dd2 100644 --- a/tensorflow/python/BUILD +++ b/tensorflow/python/BUILD @@ -222,7 +222,6 @@ py_library( "//tensorflow/python/distribute:combinations", # For tf.__internal__ API. "//tensorflow/python/distribute:distribute_config", "//tensorflow/python/distribute:estimator_training", - "//tensorflow/python/distribute:strategy_combinations", # For tf.__internal__, "//tensorflow/python/dlpack", "//tensorflow/python/eager:def_function", "//tensorflow/python/eager:monitoring", diff --git a/tensorflow/python/distribute/strategy_combinations.py b/tensorflow/python/distribute/strategy_combinations.py index 267411a61f9..57273e9bb15 100644 --- a/tensorflow/python/distribute/strategy_combinations.py +++ b/tensorflow/python/distribute/strategy_combinations.py @@ -35,9 +35,8 @@ from tensorflow.python.framework import config from tensorflow.python.platform import flags from tensorflow.python.tpu import device_assignment as device_assignment_lib from tensorflow.python.tpu import tpu_strategy_util -from tensorflow.python.util.tf_export import tf_export -_TF_INTERNAL_API_PREFIX = "__internal__.distribute.combinations." +FLAGS = flags.FLAGS _did_connect_to_cluster = False CollectiveAllReduceExtended = ( @@ -51,7 +50,6 @@ def _get_tpu_strategy_creator(steps_per_run, **kwargs): def _create_tpu_strategy(): - FLAGS = flags.FLAGS # pylint: disable=invalid-name global _did_connect_to_cluster try: @@ -196,11 +194,6 @@ mirrored_strategy_with_two_gpus = combinations.NamedDistribution( # Should call set_virtual_cpus_to_at_least(3) in your test's setUp methods. mirrored_strategy_with_cpu_1_and_2 = combinations.NamedDistribution( "Mirrored2CPU", lambda: mirrored_lib.MirroredStrategy(["/cpu:1", "/cpu:2"])) -mirrored_strategy_with_cpu_1_and_2.__doc__ = ( - """Mirrored strategy with 2 virtual CPUs. - - Should call set_virtual_cpus_to_at_least(3) in the test's setUp methods. - """) central_storage_strategy_with_two_gpus = combinations.NamedDistribution( "CentralStorage2GPUs", lambda: central_storage_strategy.CentralStorageStrategy._from_num_gpus(2), # pylint: disable=protected-access @@ -359,57 +352,3 @@ def all_strategy_minus_default_and_tpu_combinations(): def all_strategy_combinations_minus_default(): return (all_strategy_minus_default_and_tpu_combinations() + tpu_strategy_combinations()) - - -tf_export( - _TF_INTERNAL_API_PREFIX + "central_storage_strategy_with_gpu_and_cpu", - v1=[]).export_constant(__name__, - "central_storage_strategy_with_gpu_and_cpu") -tf_export( - _TF_INTERNAL_API_PREFIX + "central_storage_strategy_with_two_gpus", - v1=[]).export_constant(__name__, "central_storage_strategy_with_two_gpus") -tf_export( - _TF_INTERNAL_API_PREFIX + "cloud_tpu_strategy", - v1=[]).export_constant(__name__, "cloud_tpu_strategy") -tf_export( - _TF_INTERNAL_API_PREFIX + "default_strategy", - v1=[]).export_constant(__name__, "default_strategy") -tf_export( - _TF_INTERNAL_API_PREFIX + "mirrored_strategy_with_cpu_1_and_2", - v1=[]).export_constant(__name__, "mirrored_strategy_with_cpu_1_and_2") -tf_export( - _TF_INTERNAL_API_PREFIX + "mirrored_strategy_with_gpu_and_cpu", - v1=[]).export_constant(__name__, "mirrored_strategy_with_gpu_and_cpu") -tf_export( - _TF_INTERNAL_API_PREFIX + "mirrored_strategy_with_one_cpu", - v1=[]).export_constant(__name__, "mirrored_strategy_with_one_cpu") -tf_export( - _TF_INTERNAL_API_PREFIX + "mirrored_strategy_with_one_gpu", - v1=[]).export_constant(__name__, "mirrored_strategy_with_one_gpu") -tf_export( - _TF_INTERNAL_API_PREFIX + "mirrored_strategy_with_two_gpus", - v1=[]).export_constant(__name__, "mirrored_strategy_with_two_gpus") -tf_export( - _TF_INTERNAL_API_PREFIX + "multi_worker_mirrored_2x1_cpu", - v1=[]).export_constant(__name__, "multi_worker_mirrored_2x1_cpu") -tf_export( - _TF_INTERNAL_API_PREFIX + "multi_worker_mirrored_2x1_gpu", - v1=[]).export_constant(__name__, "multi_worker_mirrored_2x1_gpu") -tf_export( - _TF_INTERNAL_API_PREFIX + "multi_worker_mirrored_2x2_gpu", - v1=[]).export_constant(__name__, "multi_worker_mirrored_2x2_gpu") -tf_export( - _TF_INTERNAL_API_PREFIX + "one_device_strategy", - v1=[]).export_constant(__name__, "one_device_strategy") -tf_export( - _TF_INTERNAL_API_PREFIX + "one_device_strategy_gpu", - v1=[]).export_constant(__name__, "one_device_strategy_gpu") -tf_export( - _TF_INTERNAL_API_PREFIX + "tpu_strategy", - v1=[]).export_constant(__name__, "tpu_strategy") -tf_export( - _TF_INTERNAL_API_PREFIX + "tpu_strategy_one_core", - v1=[]).export_constant(__name__, "tpu_strategy_one_core") -tf_export( - _TF_INTERNAL_API_PREFIX + "tpu_strategy_packed_var", - v1=[]).export_constant(__name__, "tpu_strategy_packed_var") diff --git a/tensorflow/python/modules_with_exports.py b/tensorflow/python/modules_with_exports.py index d0dc1de542d..47b205cea14 100644 --- a/tensorflow/python/modules_with_exports.py +++ b/tensorflow/python/modules_with_exports.py @@ -55,7 +55,6 @@ from tensorflow.python.util.tf_export import tf_export # _internal APIs from tensorflow.python.distribute.combinations import generate -from tensorflow.python.distribute.strategy_combinations import * from tensorflow.python.framework.combinations import * from tensorflow.python.framework.composite_tensor import * from tensorflow.python.framework.test_combinations import * diff --git a/tensorflow/tools/api/golden/v2/tensorflow.__internal__.distribute.combinations.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.__internal__.distribute.combinations.pbtxt index c31afabba60..fe3028cb680 100644 --- a/tensorflow/tools/api/golden/v2/tensorflow.__internal__.distribute.combinations.pbtxt +++ b/tensorflow/tools/api/golden/v2/tensorflow.__internal__.distribute.combinations.pbtxt @@ -1,73 +1,5 @@ path: "tensorflow.__internal__.distribute.combinations" tf_module { - member { - name: "central_storage_strategy_with_gpu_and_cpu" - mtype: "" - } - member { - name: "central_storage_strategy_with_two_gpus" - mtype: "" - } - member { - name: "cloud_tpu_strategy" - mtype: "" - } - member { - name: "default_strategy" - mtype: "" - } - member { - name: "mirrored_strategy_with_cpu_1_and_2" - mtype: "" - } - member { - name: "mirrored_strategy_with_gpu_and_cpu" - mtype: "" - } - member { - name: "mirrored_strategy_with_one_cpu" - mtype: "" - } - member { - name: "mirrored_strategy_with_one_gpu" - mtype: "" - } - member { - name: "mirrored_strategy_with_two_gpus" - mtype: "" - } - member { - name: "multi_worker_mirrored_2x1_cpu" - mtype: "" - } - member { - name: "multi_worker_mirrored_2x1_gpu" - mtype: "" - } - member { - name: "multi_worker_mirrored_2x2_gpu" - mtype: "" - } - member { - name: "one_device_strategy" - mtype: "" - } - member { - name: "one_device_strategy_gpu" - mtype: "" - } - member { - name: "tpu_strategy" - mtype: "" - } - member { - name: "tpu_strategy_one_core" - mtype: "" - } - member { - name: "tpu_strategy_packed_var" - mtype: "" - } member_method { name: "generate" argspec: "args=[\'combinations\', \'test_combinations\'], varargs=None, keywords=None, defaults=[\'()\'], "