Keras 3 API 文档 / 多设备分布 / 分布式实用程序

分布式实用程序

[源代码]

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 号。该值应介于 0num_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