分布式工具
set_distribution function
keras.distribution.set_distribution(value)
设置为全局分布设置.
参数:
value: 一个 Distribution 实例.
distribution function
keras.distribution.distribution()
从全局上下文中检索当前分布.
list_devices function
keras.distribution.list_devices(device_type=None)
返回基于设备类型的所有可用设备.
注意:在分布式设置中,返回全局设备.
参数:
device_type:字符串,可以是 "cpu"、"gpu" 或 "tpu".
如果未提供 device_type,则默认为 "gpu" 或 "tpu"（如果可用）.
否则将返回 "cpu" 设备.
返回: 可用于分布式计算的设备列表.
initialize function
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.