Ray Tune 自定义训练函数中的检查点机制tune.report 与 Checkpoint 的完整实践【免费下载链接】rayRay is an AI compute engine. Ray consists of a core distributed runtime and a set of AI Libraries for accelerating ML workloads.项目地址: https://gitcode.com/gh_mirrors/ra/ray本文以 Ray 仓库中的示例脚本 custom_func_checkpointing.py由文档片段 custom_func_checkpointing.rst 通过literalinclude直接引入为主线讲解如何在不使用 PyTorch/TensorFlow 等官方集成、仅使用纯 Python 自定义训练函数的情况下通过ray.tune.report(metrics, checkpoint...)API 实现断点保存与恢复resume并结合 trainable_fn_utils.py 与 Checkpoint 类实现 说明其底层机制。读完后你将掌握自定义训练函数中检查点的写入与读取方式、Checkpoint对象的目录语义以及Tuner层面超参搜索配置的完整可运行代码。核心思路用tune.report(checkpoint...)报告检查点当训练逻辑是自定义函数而非 Ray 的 PyTorch/TF 等集成时Tune 无法替你感知训练状态因此需要你主动做两件事写入在训练循环中周期性地把状态如当前 step、模型权重序列化到一个目录中然后通过tune.report(..., checkpointCheckpoint.from_directory(...))交给 Tune 持久化恢复在训练函数入口处调用tune.get_checkpoint()获取最近一次已报告的检查点反序列化后从断点继续。示例完整代码如下来自 custom_func_checkpointing.pyimport argparse import json import os import tempfile import time from ray import tune from ray.tune import Checkpoint def evaluation_fn(step, width, height): time.sleep(0.1) return (0.1 width * step / 100) ** (-1) height * 0.1 def train_func(config): step 0 width, height config[width], config[height] checkpoint tune.get_checkpoint() if checkpoint: with checkpoint.as_directory() as checkpoint_dir: with open(os.path.join(checkpoint_dir, checkpoint.json)) as f: state json.load(f) step state[step] 1 for current_step in range(step, 100): intermediate_score evaluation_fn(current_step, width, height) with tempfile.TemporaryDirectory() as temp_checkpoint_dir: with open(os.path.join(temp_checkpoint_dir, checkpoint.json), w) as f: json.dump({step: current_step}, f) tune.report( {iterations: current_step, mean_loss: intermediate_score}, checkpointCheckpoint.from_directory(temp_checkpoint_dir), ) if __name__ __main__: parser argparse.ArgumentParser() parser.add_argument( --smoke-test, actionstore_true, helpFinish quickly for testing ) args, _ parser.parse_known_args() tuner tune.Tuner( train_func, run_configtune.RunConfig( namehyperband_test, stop{training_iteration: 1 if args.smoke_test else 10}, ), tune_configtune.TuneConfig( metricmean_loss, modemin, num_samples5, ), param_space{ steps: 10, width: tune.randint(10, 100), height: tune.loguniform(10, 100), }, ) results tuner.fit() best_result results.get_best_result() print(Best hyperparameters: , best_result.config) best_checkpoint best_result.checkpoint print(Best checkpoint: , best_checkpoint)脚本支持--smoke-test命令行参数传入后stop中的training_iteration从默认的 10 降为 1便于快速验证流水线。train_func的恢复逻辑tune.get_checkpoint()的语义训练函数开头的关键代码是checkpoint tune.get_checkpoint() if checkpoint: with checkpoint.as_directory() as checkpoint_dir: with open(os.path.join(checkpoint_dir, checkpoint.json)) as f: state json.load(f) step state[step] 1从源码结构看trainable_fn_utils.py 中get_checkpoint()的实现是PublicAPI(stabilitystable) _warn_session_misuse() def get_checkpoint() - Optional[Checkpoint]: Access the latest reported checkpoint to resume from if one exists. return get_session().loaded_checkpoint即它返回的是当前 Train 会话session中最近加载的检查点首次运行时没有任何历史检查点返回None训练函数从step 0开始。当实验因故障恢复或断点续跑时Tune 会把上一次报告的检查点装载进会话——这一点在 function_trainable.py 中可以看到闭环session.loaded_checkpoint checkpoint_result.checkpoint也就是说每一次tune.report(..., checkpoint...)的结果都会写回session.loaded_checkpoint形成报告 → 持久化 → 恢复时加载的完整链路。as_directory()把检查点当作只读本地目录Checkpoint.as_directory()是上下文管理器。从 Checkpoint 实现 的文档说明可以确认其两种行为本地目录检查点直接返回原目录路径不做拷贝退出上下文后不做清理远端存储检查点如 S3 URI下载到本地临时目录后返回路径退出上下文时清理临时目录若同一节点上多个进程并发访问同一检查点只有一个进程真正执行下载其余进程等待后共享同一份数据通过TempFileLock文件锁实现。官方文档同时强调返回的目录应视为只读因为临时数据可能在退出上下文后被删除。示例中读取checkpoint.json恢复 step 的做法正是这种目录即状态文件集合的典型用法——检查点的物理形态就是一个目录用户完全自由决定其中放什么JSON、pickle、模型权重文件等。写入检查点tune.reportCheckpoint.from_directory训练循环中每步都执行一次检查点报告with tempfile.TemporaryDirectory() as temp_checkpoint_dir: with open(os.path.join(temp_checkpoint_dir, checkpoint.json), w) as f: json.dump({step: current_step}, f) tune.report( {iterations: current_step, mean_loss: intermediate_score}, checkpointCheckpoint.from_directory(temp_checkpoint_dir), )这里有两个值得注意的细节先写临时目录再包装成Checkpoint对象Checkpoint.from_directory(path)源码会以pyarrow.fs.LocalFileSystem()为后端构造Checkpoint把本地目录变成可被存储层持久化/迁移的检查点引用。tune.report的 docstring 明确说明提供checkpoint时它会被持久化到配置的存储位置persistent storage。metrics 与 checkpoint 在同一次调用中报告tune.report的 源码 表明每次调用都会自动递增底层的training_iteration计数——示例main块中stop{training_iteration: 10}正是以报告次数为停止条件而非 epoch 数。docstring 也提醒这个iteration的物理含义由用户按调用report的频率自行定义不一定对应一个 epoch。另外注意 trainable_fn_utils.py 中的实现细节_copy_doc(TrainCheckpoint) class Checkpoint(TrainCheckpoint): # NOTE: This is just a pass-through wrapper around ray.train.Checkpoint # in order to detect whether the import module was correct ray.tune.Checkpoint. passray.tune.Checkpoint只是ray.train.Checkpoint的透传包装类专门用于检测导入模块是否正确若你在report中传入了错误来源的Checkpoint实例会触发 v2 迁移弃用警告见 report 函数中的类型检查因此务必从ray.tune导入Checkpoint与示例文件开头的from ray.tune import Checkpoint保持一致。Checkpoint类的更多能力示例未用到但实用从 Checkpoint 完整实现 可以看到该类还支持远端构造Checkpoint(s3://bucket/path/to/checkpoint)会根据 URI scheme 推断文件系统无需显式传入 filesystemto_directory(pathNone)把检查点内容写到指定本地目录或自动生成的临时目录适用于需要持久保留下载内容的场景元数据 APIset_metadata/get_metadata/update_metadata会把键值对以.metadata.json形式随检查点持久化适合存放超参摘要、预处理器配置等误用防护__fspath__被刻意实现为抛TypeError强制你使用to_directory()/as_directory()而非把Checkpoint当普通路径拼接使用。Tuner配置逐项解析示例main块展示了完整的tune.Tuner调用各部分含义如下配置项取值作用train_func位置参数上面的训练函数Tune 自动包装为 Function TrainableRunConfig.namehyperband_test实验名用于标识运行目录RunConfig.stop{training_iteration: 1}smoke/10默认每个 trial 报告满指定 iteration 次数即停止TuneConfig.metricmean_loss用于排序与早停的目标指标每次report中提供TuneConfig.modemin目标越小越好loss 场景TuneConfig.num_samples5每个 trial 重复采样 5 次以抑制单次随机性param_space[steps]10常量非随机超参仅随配置下发示例中实际未使用param_space[width]tune.randint(10, 100)在 [10, 100] 均匀取整param_space[height]tune.loguniform(10, 100)在 [10, 100] 对数均匀采样适合量级跨度大的参数训练结束后通过results.get_best_result()取最优 trial并打印其config与checkpoint——best_result.checkpoint即最后一次报告的Checkpoint对象可直接用as_directory()加载其中的checkpoint.json或真实场景中的模型权重用于部署/推理。运行方式与适用前提python python/ray/tune/examples/custom_func_checkpointing.py # 完整跑 10 次 iteration python python/ray/tune/examples/custom_func_checkpointing.py --smoke-test # 1 次 iteration 快速验证适用前提与限制训练函数运行在 Ray 集群的 worker 中tune.report/tune.get_checkpoint只能在 Train 会话内调用源码中由_warn_session_misuse()装饰器检测误用检查点内容本身只是目录数据Tune 不负责解释其结构恢复逻辑如step 1完全由用户编写该模式与官方集成如 cifar10_pytorch.py、pbt_convnet_function_example.py 中对tune.get_checkpoint()的同款用法遵循同一套 API因此从自定义函数迁移到框架集成时检查点代码基本可以平移。小结Ray Tune 的自定义函数检查点机制可归纳为三步闭环用Checkpoint.from_directory把任意本地目录包装为检查点 → 随tune.report一起上报以触发持久化 → 恢复时经tune.get_checkpoint().as_directory()以只读目录形式读回。整套机制不依赖任何深度学习框架检查点可以是 JSON、pickle 或模型权重文件的任意组合而training_iteration随每次report自动递增的特性使得RunConfig.stop可以按报告次数精确控制试验时长。【免费下载链接】rayRay is an AI compute engine. Ray consists of a core distributed runtime and a set of AI Libraries for accelerating ML workloads.项目地址: https://gitcode.com/gh_mirrors/ra/ray创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 Ray Tune 自定义训练函数中的检查点机制：tune.report 与 Checkpoint 的完整实践
Ray Tune 自定义训练函数中的检查点机制tune.report 与 Checkpoint 的完整实践【免费下载链接】rayRay is an AI compute engine. Ray consists of a core distributed runtime and a set of AI Libraries for accelerating ML workloads.项目地址: https://gitcode.com/gh_mirrors/ra/ray本文以 Ray 仓库中的示例脚本 custom_func_checkpointing.py由文档片段 custom_func_checkpointing.rst 通过literalinclude直接引入为主线讲解如何在不使用 PyTorch/TensorFlow 等官方集成、仅使用纯 Python 自定义训练函数的情况下通过ray.tune.report(metrics, checkpoint...)API 实现断点保存与恢复resume并结合 trainable_fn_utils.py 与 Checkpoint 类实现 说明其底层机制。读完后你将掌握自定义训练函数中检查点的写入与读取方式、Checkpoint对象的目录语义以及Tuner层面超参搜索配置的完整可运行代码。核心思路用tune.report(checkpoint...)报告检查点当训练逻辑是自定义函数而非 Ray 的 PyTorch/TF 等集成时Tune 无法替你感知训练状态因此需要你主动做两件事写入在训练循环中周期性地把状态如当前 step、模型权重序列化到一个目录中然后通过tune.report(..., checkpointCheckpoint.from_directory(...))交给 Tune 持久化恢复在训练函数入口处调用tune.get_checkpoint()获取最近一次已报告的检查点反序列化后从断点继续。示例完整代码如下来自 custom_func_checkpointing.pyimport argparse import json import os import tempfile import time from ray import tune from ray.tune import Checkpoint def evaluation_fn(step, width, height): time.sleep(0.1) return (0.1 width * step / 100) ** (-1) height * 0.1 def train_func(config): step 0 width, height config[width], config[height] checkpoint tune.get_checkpoint() if checkpoint: with checkpoint.as_directory() as checkpoint_dir: with open(os.path.join(checkpoint_dir, checkpoint.json)) as f: state json.load(f) step state[step] 1 for current_step in range(step, 100): intermediate_score evaluation_fn(current_step, width, height) with tempfile.TemporaryDirectory() as temp_checkpoint_dir: with open(os.path.join(temp_checkpoint_dir, checkpoint.json), w) as f: json.dump({step: current_step}, f) tune.report( {iterations: current_step, mean_loss: intermediate_score}, checkpointCheckpoint.from_directory(temp_checkpoint_dir), ) if __name__ __main__: parser argparse.ArgumentParser() parser.add_argument( --smoke-test, actionstore_true, helpFinish quickly for testing ) args, _ parser.parse_known_args() tuner tune.Tuner( train_func, run_configtune.RunConfig( namehyperband_test, stop{training_iteration: 1 if args.smoke_test else 10}, ), tune_configtune.TuneConfig( metricmean_loss, modemin, num_samples5, ), param_space{ steps: 10, width: tune.randint(10, 100), height: tune.loguniform(10, 100), }, ) results tuner.fit() best_result results.get_best_result() print(Best hyperparameters: , best_result.config) best_checkpoint best_result.checkpoint print(Best checkpoint: , best_checkpoint)脚本支持--smoke-test命令行参数传入后stop中的training_iteration从默认的 10 降为 1便于快速验证流水线。train_func的恢复逻辑tune.get_checkpoint()的语义训练函数开头的关键代码是checkpoint tune.get_checkpoint() if checkpoint: with checkpoint.as_directory() as checkpoint_dir: with open(os.path.join(checkpoint_dir, checkpoint.json)) as f: state json.load(f) step state[step] 1从源码结构看trainable_fn_utils.py 中get_checkpoint()的实现是PublicAPI(stabilitystable) _warn_session_misuse() def get_checkpoint() - Optional[Checkpoint]: Access the latest reported checkpoint to resume from if one exists. return get_session().loaded_checkpoint即它返回的是当前 Train 会话session中最近加载的检查点首次运行时没有任何历史检查点返回None训练函数从step 0开始。当实验因故障恢复或断点续跑时Tune 会把上一次报告的检查点装载进会话——这一点在 function_trainable.py 中可以看到闭环session.loaded_checkpoint checkpoint_result.checkpoint也就是说每一次tune.report(..., checkpoint...)的结果都会写回session.loaded_checkpoint形成报告 → 持久化 → 恢复时加载的完整链路。as_directory()把检查点当作只读本地目录Checkpoint.as_directory()是上下文管理器。从 Checkpoint 实现 的文档说明可以确认其两种行为本地目录检查点直接返回原目录路径不做拷贝退出上下文后不做清理远端存储检查点如 S3 URI下载到本地临时目录后返回路径退出上下文时清理临时目录若同一节点上多个进程并发访问同一检查点只有一个进程真正执行下载其余进程等待后共享同一份数据通过TempFileLock文件锁实现。官方文档同时强调返回的目录应视为只读因为临时数据可能在退出上下文后被删除。示例中读取checkpoint.json恢复 step 的做法正是这种目录即状态文件集合的典型用法——检查点的物理形态就是一个目录用户完全自由决定其中放什么JSON、pickle、模型权重文件等。写入检查点tune.reportCheckpoint.from_directory训练循环中每步都执行一次检查点报告with tempfile.TemporaryDirectory() as temp_checkpoint_dir: with open(os.path.join(temp_checkpoint_dir, checkpoint.json), w) as f: json.dump({step: current_step}, f) tune.report( {iterations: current_step, mean_loss: intermediate_score}, checkpointCheckpoint.from_directory(temp_checkpoint_dir), )这里有两个值得注意的细节先写临时目录再包装成Checkpoint对象Checkpoint.from_directory(path)源码会以pyarrow.fs.LocalFileSystem()为后端构造Checkpoint把本地目录变成可被存储层持久化/迁移的检查点引用。tune.report的 docstring 明确说明提供checkpoint时它会被持久化到配置的存储位置persistent storage。metrics 与 checkpoint 在同一次调用中报告tune.report的 源码 表明每次调用都会自动递增底层的training_iteration计数——示例main块中stop{training_iteration: 10}正是以报告次数为停止条件而非 epoch 数。docstring 也提醒这个iteration的物理含义由用户按调用report的频率自行定义不一定对应一个 epoch。另外注意 trainable_fn_utils.py 中的实现细节_copy_doc(TrainCheckpoint) class Checkpoint(TrainCheckpoint): # NOTE: This is just a pass-through wrapper around ray.train.Checkpoint # in order to detect whether the import module was correct ray.tune.Checkpoint. passray.tune.Checkpoint只是ray.train.Checkpoint的透传包装类专门用于检测导入模块是否正确若你在report中传入了错误来源的Checkpoint实例会触发 v2 迁移弃用警告见 report 函数中的类型检查因此务必从ray.tune导入Checkpoint与示例文件开头的from ray.tune import Checkpoint保持一致。Checkpoint类的更多能力示例未用到但实用从 Checkpoint 完整实现 可以看到该类还支持远端构造Checkpoint(s3://bucket/path/to/checkpoint)会根据 URI scheme 推断文件系统无需显式传入 filesystemto_directory(pathNone)把检查点内容写到指定本地目录或自动生成的临时目录适用于需要持久保留下载内容的场景元数据 APIset_metadata/get_metadata/update_metadata会把键值对以.metadata.json形式随检查点持久化适合存放超参摘要、预处理器配置等误用防护__fspath__被刻意实现为抛TypeError强制你使用to_directory()/as_directory()而非把Checkpoint当普通路径拼接使用。Tuner配置逐项解析示例main块展示了完整的tune.Tuner调用各部分含义如下配置项取值作用train_func位置参数上面的训练函数Tune 自动包装为 Function TrainableRunConfig.namehyperband_test实验名用于标识运行目录RunConfig.stop{training_iteration: 1}smoke/10默认每个 trial 报告满指定 iteration 次数即停止TuneConfig.metricmean_loss用于排序与早停的目标指标每次report中提供TuneConfig.modemin目标越小越好loss 场景TuneConfig.num_samples5每个 trial 重复采样 5 次以抑制单次随机性param_space[steps]10常量非随机超参仅随配置下发示例中实际未使用param_space[width]tune.randint(10, 100)在 [10, 100] 均匀取整param_space[height]tune.loguniform(10, 100)在 [10, 100] 对数均匀采样适合量级跨度大的参数训练结束后通过results.get_best_result()取最优 trial并打印其config与checkpoint——best_result.checkpoint即最后一次报告的Checkpoint对象可直接用as_directory()加载其中的checkpoint.json或真实场景中的模型权重用于部署/推理。运行方式与适用前提python python/ray/tune/examples/custom_func_checkpointing.py # 完整跑 10 次 iteration python python/ray/tune/examples/custom_func_checkpointing.py --smoke-test # 1 次 iteration 快速验证适用前提与限制训练函数运行在 Ray 集群的 worker 中tune.report/tune.get_checkpoint只能在 Train 会话内调用源码中由_warn_session_misuse()装饰器检测误用检查点内容本身只是目录数据Tune 不负责解释其结构恢复逻辑如step 1完全由用户编写该模式与官方集成如 cifar10_pytorch.py、pbt_convnet_function_example.py 中对tune.get_checkpoint()的同款用法遵循同一套 API因此从自定义函数迁移到框架集成时检查点代码基本可以平移。小结Ray Tune 的自定义函数检查点机制可归纳为三步闭环用Checkpoint.from_directory把任意本地目录包装为检查点 → 随tune.report一起上报以触发持久化 → 恢复时经tune.get_checkpoint().as_directory()以只读目录形式读回。整套机制不依赖任何深度学习框架检查点可以是 JSON、pickle 或模型权重文件的任意组合而training_iteration随每次report自动递增的特性使得RunConfig.stop可以按报告次数精确控制试验时长。【免费下载链接】rayRay is an AI compute engine. Ray consists of a core distributed runtime and a set of AI Libraries for accelerating ML workloads.项目地址: https://gitcode.com/gh_mirrors/ra/ray创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考