分发工具

[源代码]

set_distribution 函数

keras.distribution.set_distribution(value)

将分发设置为全局分发配置。

参数

• value：一个 Distribution 实例。

[源代码]

distribution 函数

keras.distribution.distribution()

从全局上下文中检索当前分发。

[源代码]

list_devices 函数

keras.distribution.list_devices(device_type=None)

根据设备类型返回所有可用设备。

注意：在分布式设置中，返回的是全局设备。

参数

• device_type：字符串，可以是 "cpu"、"gpu" 或 "tpu" 之一。如果未提供 device_type，则默认为可用的 "gpu" 或 "tpu"。否则将返回 "cpu" 设备。

返回值：可用于分布式计算的设备列表。

[源代码]

initialize 函数

keras.distribution.initialize(
    job_addresses=None, num_processes=None, process_id=None
)

为多主机/进程设置初始化分发系统。

调用 initialize 将为在多主机 GPU 或 TPU 上执行准备后端。它应在任何计算之前调用。

请注意，参数也可以通过环境变量注入，这可以通过启动脚本在启动时更好地控制。对于某些也依赖环境变量进行配置的后端，Keras 将正确地转发它们。

参数

• job_addresses：字符串。组成整个计算集群的所有作业的 IP 地址，用逗号分隔。请注意，对于 JAX 后端，只需要作业 0（协调器）的地址。对于某些运行时，例如云 TPU，此值可以为 None，后端将通过 TPU 环境变量确定它。您也可以通过环境变量 KERAS_DISTRIBUTION_JOB_ADDRESSES 配置此值。
• num_processes：整数。组成整个计算集群的工作进程/进程数。对于某些运行时，例如云 TPU，此值可以为 None，后端将通过 TPU 环境变量确定它。您也可以通过环境变量 KERAS_DISTRIBUTION_NUM_PROCESSES 配置此值。
• process_id：整数。当前工作进程/进程的 ID 号。此值应在 0 到 num_processes - 1 范围内。0 表示当前工作进程/进程是主/协调作业。您也可以通过环境变量 KERAS_DISTRIBUTION_PROCESS_ID 配置此值。

示例

假设有两个 GPU 进程，进程 0 运行在地址 10.0.0.1:1234，进程 1 运行在地址 10.0.0.2:2345。要配置此类集群，您可以运行 - __ 在进程 0 上运行__

keras.distribute.initialize(
  job_addresses="10.0.0.1:1234,10.0.0.2:2345",
  num_processes=2,
  process_id=0)

• __ 在进程 1 上运行__

keras.distribute.initialize(
  job_addresses="10.0.0.1:1234,10.0.0.2:2345",
  num_processes=2,
  process_id=1)

• __ 或通过环境变量__
• 在进程 0 上:

os.environ[
  "KERAS_DISTRIBUTION_JOB_ADDRESSES"] = "10.0.0.1:1234,10.0.0.2:2345"
os.environ["KERAS_DISTRIBUTION_NUM_PROCESSES"] = "2"
os.environ["KERAS_DISTRIBUTION_PROCESS_ID"] = "0"
keras.distribute.initialize()

• __ 在进程 1 上运行__

os.environ[
  "KERAS_DISTRIBUTION_JOB_ADDRESSES"] = "10.0.0.1:1234,10.0.0.2:2345"
os.environ["KERAS_DISTRIBUTION_NUM_PROCESSES"] = "2"
os.environ["KERAS_DISTRIBUTION_PROCESS_ID"] = "1"
keras.distribute.initialize()

另请注意，对于 JAX 后端，job_addresses 可以进一步简化为仅包含主/协调器地址，即 - __10.0.0.1__:1234。