69 lines
3.0 KiB
Python
69 lines
3.0 KiB
Python
|
# 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.
|
||
|
|
# ==============================================================================
|
||
|
|
"""Strategy combinations for combinations.combine()."""
|
||
|
|
|
||
|
|
import tensorflow.compat.v2 as tf
|
||
|
|
|
||
|
|
multidevice_strategies = [
|
||
|
|
tf.__internal__.distribute.combinations.mirrored_strategy_with_gpu_and_cpu,
|
||
|
|
tf.__internal__.distribute.combinations.mirrored_strategy_with_two_gpus,
|
||
|
|
tf.__internal__.distribute.combinations.tpu_strategy,
|
||
|
|
]
|
||
|
|
|
||
|
|
multiworker_strategies = [
|
||
|
|
tf.__internal__.distribute.combinations.multi_worker_mirrored_2x1_cpu,
|
||
|
|
tf.__internal__.distribute.combinations.multi_worker_mirrored_2x1_gpu,
|
||
|
|
tf.__internal__.distribute.combinations.multi_worker_mirrored_2x2_gpu,
|
||
|
|
]
|
||
|
|
|
||
|
|
strategies_minus_default_minus_tpu = [
|
||
|
|
tf.__internal__.distribute.combinations.one_device_strategy,
|
||
|
|
tf.__internal__.distribute.combinations.one_device_strategy_gpu,
|
||
|
|
tf.__internal__.distribute.combinations.mirrored_strategy_with_gpu_and_cpu,
|
||
|
|
tf.__internal__.distribute.combinations.mirrored_strategy_with_two_gpus,
|
||
|
|
tf.__internal__.distribute.combinations.central_storage_strategy_with_gpu_and_cpu, # noqa: E501
|
||
|
|
]
|
||
|
|
|
||
|
|
strategies_minus_tpu = [
|
||
|
|
tf.__internal__.distribute.combinations.default_strategy,
|
||
|
|
tf.__internal__.distribute.combinations.one_device_strategy,
|
||
|
|
tf.__internal__.distribute.combinations.one_device_strategy_gpu,
|
||
|
|
tf.__internal__.distribute.combinations.mirrored_strategy_with_gpu_and_cpu,
|
||
|
|
tf.__internal__.distribute.combinations.mirrored_strategy_with_two_gpus,
|
||
|
|
tf.__internal__.distribute.combinations.central_storage_strategy_with_gpu_and_cpu, # noqa: E501
|
||
|
|
]
|
||
|
|
|
||
|
|
multi_worker_mirrored_strategies = [
|
||
|
|
tf.__internal__.distribute.combinations.multi_worker_mirrored_2x1_cpu,
|
||
|
|
tf.__internal__.distribute.combinations.multi_worker_mirrored_2x1_gpu,
|
||
|
|
tf.__internal__.distribute.combinations.multi_worker_mirrored_2x2_gpu,
|
||
|
|
]
|
||
|
|
|
||
|
|
tpu_strategies = [
|
||
|
|
tf.__internal__.distribute.combinations.tpu_strategy,
|
||
|
|
]
|
||
|
|
|
||
|
|
parameter_server_strategies_single_worker = [
|
||
|
|
tf.__internal__.distribute.combinations.parameter_server_strategy_1worker_2ps_cpu, # noqa: E501
|
||
|
|
tf.__internal__.distribute.combinations.parameter_server_strategy_1worker_2ps_1gpu, # noqa: E501
|
||
|
|
]
|
||
|
|
|
||
|
|
parameter_server_strategies_multi_worker = [
|
||
|
|
tf.__internal__.distribute.combinations.parameter_server_strategy_3worker_2ps_cpu, # noqa: E501
|
||
|
|
tf.__internal__.distribute.combinations.parameter_server_strategy_3worker_2ps_1gpu, # noqa: E501
|
||
|
|
]
|
||
|
|
|
||
|
|
all_strategies = strategies_minus_tpu + tpu_strategies
|