Determined AI Core API 入门指南：从基础到分布式训练-CSDN博客

本文链接：https://blog.csdn.net/gitblog_00513/article/details/148602402

Determined AI Core API 入门指南：从基础到分布式训练

determined Determined is an open-source machine learning platform that simplifies distributed training, hyperparameter tuning, experiment tracking, and resource management. Works with PyTorch and TensorFlow. 项目地址: https://gitcode.com/gh_mirrors/de/determined

概述

Determined AI 是一个强大的机器学习平台，其 Core API 为开发者提供了灵活的方式来构建和训练模型。本文将带你从零开始学习如何使用 Core API，通过一个简单的整数递增示例，逐步掌握核心功能。

Core API 核心优势

Core API 作为 Determined 平台的基础 API，具有以下显著优势：

指标追踪：轻松记录训练和验证过程中的各项指标
检查点管理：支持模型检查点的保存和恢复
超参数搜索：内置高级超参数优化功能
分布式训练：无缝支持多 GPU 和多节点训练

基础示例：整数递增训练

我们从最简单的示例开始 - 一个循环递增整数的"模型"：

# 0_start.py
import logging

def main():
    x = 0
    max_length = 100
    for batch in range(max_length):
        x += 1
        logging.info(f"x is now {x}")

if __name__ == "__main__":
    main()

对应的配置文件 0_start.yaml 只需包含基本配置：

name: core_api_example
entrypoint: python3 0_start.py

指标报告功能

在实际应用中，我们需要向平台报告训练指标。以下是关键步骤：

导入 Determined 核心模块
配置日志格式
使用 core.init 初始化上下文
通过上下文对象报告指标

# 1_metrics.py 关键部分
import determined as det

def main(core_context):
    x = 0
    max_length = 100
    for batch in range(max_length):
        x += 1
        # 每10步报告训练指标
        if batch % 10 == 0:
            core_context.train.report_training_metrics(
                steps_completed=batch,
                metrics={"x": x},
            )
        # 模拟验证过程
        if batch % 20 == 0:
            core_context.train.report_validation_metrics(
                steps_completed=batch,
                metrics={"val_x": x},
            )

检查点管理

实现检查点功能需要考虑两种场景：

暂停后恢复训练（保持批次索引）
全新继续训练（重置批次索引）

关键实现要点：

def save_state(core_context, x, batch, trial_id):
    with core_context.checkpoint.store_path({"model_data": None}) as path:
        state = {"x": x, "batch": batch, "trial_id": trial_id}
        torch.save(state, path / "state.pth")

def load_state(checkpoint):
    with checkpoint.restore_path() as path:
        return torch.load(path / "state.pth")

超参数搜索

Core API 支持高级超参数搜索策略：

# 3_hpsearch.yaml 配置示例
hyperparameters:
  increment_by:
    type: int
    minval: 1
    maxval: 5

在训练脚本中获取超参数：

hparams = info.trial.hparams
x += hparams["increment_by"]

分布式训练实现

分布式训练需要特殊处理：

使用 allgather 等通信原语
主工作节点管理
同步预emption检查

关键代码结构：

def worker_main(rank, core_context, increment_by):
    # 分布式训练逻辑
    all_increment_bys = core_context.distributed.allgather(increment_by)
    x += sum(all_increment_bys)
    
    # 仅主节点报告指标
    if core_context.distributed.rank == 0:
        core_context.train.report_training_metrics(...)