企业官网建设流程全解析

NNI Trial 开发指南：如何编写运行在 NNI 上的 Trial（Tuner 参数获取与结果上报全流程）

【免费下载链接】nniAn open source AutoML toolkit for automate machine learning lifecycle, including feature engineering, neural architecture search, model compression and hyper-parameter tuning.项目地址: https://gitcode.com/gh_mirrors/nn/nni

Trial 是 NNI（Neural Network Intelligence）自动机器学习框架中执行训练任务的最小单元：它从 Tuner 接收超参数/网络结构配置，将训练过程中的中间结果发送给 Assessor，并将最终结果回传给 Tuner。本文基于仓库中的 examples/trials/README.md 及对应的 mnist-keras 完整示例，系统讲解如何把一段普通的机器学习代码改造成可运行在 NNI 上的 Trial。读完本文，你将掌握 Trial 与 NNI 框架交互的四个核心步骤：准备可运行的原始代码、通过nni.get_next_parameter()获取配置、通过nni.report_intermediate_result()上报中间结果、通过nni.report_final_result()上报最终结果，并能独立搭建一个可被 Tuner 调参、可被 Assessor 提前终止的完整实验。

Trial 在 NNI 中的角色与职责

在 NNI 的自动机器学习生命周期中，Trial 是承上启下的执行节点。仓库文档开篇即点明其定位：

Trial receive the hyper-parameter/architecture configure from Tuner, and send intermediate result to Assessor and final result to Tuner.

即：Trial 的输入是 Tuner 根据搜索空间（search space）采样出的参数，输出是训练过程中产生的中间指标与训练结束后的最终指标。这一交互在源码层面有明确对应。在 nni/trial.py 中，Trial 侧暴露的全部 API 为：

• get_next_parameter()/get_next_parameters()：获取 Tuner 生成的超参数；
• get_current_parameter(tag)：读取当前参数（可带字段名）；
• report_intermediate_result(metric)：上报中间结果；
• report_final_result(metric)：上报最终结果；
• get_experiment_id()/get_trial_id()/get_sequence_id()：获取实验 ID、Trial ID 与序号。

因此，编写一个运行在 NNI 上的 Trial 通常只需要四步：先有一份能在本地跑通的机器学习代码，再按本文所述插入 NNI 的 API 调用。

第一步：准备一份可运行的原始 Trial

Trial 的本质是一段能够在本地直接运行的机器学习代码，NNI 不要求对模型结构做任何特殊改写。文档以mnist-keras.py为例展示了一份"原始"代码：使用 Keras 构建一个简单的卷积网络（两层Conv2D+MaxPooling2D+Flatten+ 两层Dense），加载 MNIST 数据集训练并评估。其关键点在于：

import argparse import logging import keras import numpy as np from keras import backend as K from keras.datasets import mnist from keras.layers import Conv2D, Dense, Flatten, MaxPooling2D from keras.models import Sequential K.set_image_data_format('channels_last') H, W = 28, 28 NUM_CLASSES = 10 def create_mnist_model(hyper_params, input_shape=(H, W, 1), num_classes=NUM_CLASSES): layers = [ Conv2D(32, kernel_size=(3, 3), activation='relu', input_shape=input_shape), Conv2D(64, (3, 3), activation='relu'), MaxPooling2D(pool_size=(2, 2)), Flatten(), Dense(100, activation='relu'), Dense(num_classes, activation='softmax') ] model = Sequential(layers) if hyper_params['optimizer'] == 'Adam': optimizer = keras.optimizers.Adam(lr=hyper_params['learning_rate']) else: optimizer = keras.optimizers.SGD(lr=hyper_params['learning_rate'], momentum=0.9) model.compile(loss=keras.losses.categorical_crossentropy, optimizer=optimizer, metrics=['accuracy']) return model def load_mnist_data(args): (x_train, y_train), (x_test, y_test) = mnist.load_data() x_train = (np.expand_dims(x_train, -1).astype(float) / 255.)[:args.num_train] x_test = (np.expand_dims(x_test, -1).astype(float) / 255.)[:args.num_test] y_train = keras.utils.to_categorical(y_train, NUM_CLASSES)[:args.num_train] y_test = keras.utils.to_categorical(y_test, NUM_CLASSES)[:args.num_test] return x_train, y_train, x_test, y_test def train(args, params): x_train, y_train, x_test, y_test = load_mnist_data(args) model = create_mnist_model(params) model.fit(x_train, y_train, batch_size=args.batch_size, epochs=args.epochs, verbose=1, validation_data=(x_test, y_test), callbacks=[SendMetrics()]) _, acc = model.evaluate(x_test, y_test, verbose=0) def generate_default_params(): return { 'optimizer': 'Adam', 'learning_rate': 0.001 } if __name__ == '__main__': PARSER = argparse.ArgumentParser() PARSER.add_argument("--batch_size", type=int, default=200, help="batch size", required=False) PARSER.add_argument("--epochs", type=int, default=10, help="Train epochs", required=False) PARSER.add_argument("--num_train", type=int, default=1000, help="Number of train samples to be used, maximum 60000", required=False) PARSER.add_argument("--num_test", type=int, default=1000, help="Number of test samples to be used, maximum 10000", required=False) ARGS, UNKNOWN = PARSER.parse_known_args() PARAMS = generate_default_params() train(ARGS, PARAMS)

这段代码没有任何 NNI 依赖，可直接在本地运行，用于验证代码正确性。仓库中的实际示例 examples/trials/mnist-keras/mnist-keras.py 在此基础上进一步做了两处工程化增强：一是通过os.environ['NNI_OUTPUT_DIR']把 TensorBoard 日志目录指向 NNI 分配的输出目录，二是将 MNIST 数据集缓存到NNI_OUTPUT_DIR下并在使用后删除，避免多 Trial 并发时相互污染。

注意其中的SendMetrics回调在原始代码中是空实现（pass），这正是后续要接入 NNI 的位置。

第二步：从 Tuner 获取超参数配置

改造的第一处关键动作是引入nni模块并调用nni.get_next_parameter()。文档特别提醒关注示例中的第 10、24、25 行，即导入语句、调用获取参数的语句以及用返回结果更新默认参数字典的语句：

import nni # 第 10 行：导入 nni if __name__ == '__main__': PARSER = argparse.ArgumentParser() ... ARGS, UNKNOWN = PARSER.parse_known_args() PARAMS = generate_default_params() RECEIVED_PARAMS = nni.get_next_parameter() # 获取 Tuner 采样出的参数 PARAMS.update(RECEIVED_PARAMS) # 用 Tuner 参数覆盖默认值 train(ARGS, PARAMS)

这一模式非常关键：先用generate_default_params()提供一份可独立运行的默认参数，再用 Tuner 返回的参数update覆盖默认值。这样既保证"无 Tuner 参数时也能跑通"，又保证参数真正来自 Tuner。

从源码看，nni.get_next_parameter()在 nni/trial.py 中通过get_default_trial_command_channel().receive_parameter()从 NNI manager 接收参数记录，并返回其中的parameters字段。其 docstring 给出了典型的返回形态：若搜索空间为{'activation': {'_type': 'choice', '_value': ['relu', 'tanh', 'sigmoid']}, 'learning_rate': {'_type': 'loguniform', '_value': [0.0001, 0.1]}}，则返回值形如{'activation': 'relu', 'learning_rate': 0.02}。同时源码明确要求："每个 Trial 应且只应调用一次该函数，否则行为未定义"（见 docstring），这是编写 Trial 时必须遵守的约定。

值得注意的是，Trial 代码在脱离 NNI 环境独立运行时，receive_parameter()会走 nni/runtime/trial_command_channel/standalone.py 中的StandaloneTrialCommandChannel，返回空参数集{}并发出运行时警告，从而保证同一份代码既能被 NNI 调度、也能本地调试。

第三步：上报中间结果给 Assessor

中间结果（intermediate result）是训练过程中周期产生的指标，典型的就是每个 epoch 的 accuracy 或 loss。它的接收方是 Assessor（如早停算法），NNI 据此决定是否提前终止表现不佳的 Trial，从而节省计算资源。

在 Keras 中，最自然的接入点是回调（Callback）。文档改造了SendMetrics回调，在on_epoch_end中调用nni.report_intermediate_result(logs)：

class SendMetrics(keras.callbacks.Callback): def on_epoch_end(self, epoch, logs={}): nni.report_intermediate_result(logs)

在model.fit(...)时把该回调传入callbacks=[SendMetrics()]即可。

仓库中的真实示例 examples/trials/mnist-keras/mnist-keras.py 做了一个值得借鉴的健壮性处理：Keras 不同版本中验证集准确率的日志键名不一致（TensorFlow 2.0 文档称其为val_acc，实际为val_accuracy），因此它同时兼容两种情况：

if 'val_acc' in logs: nni.report_intermediate_result(logs['val_acc']) else: nni.report_intermediate_result(logs['val_accuracy'])

这提示了一个通用原则：上报的指标值应确保是框架期望的数值形态。

第四步：上报最终结果给 Tuner

训练结束后，Trial 需要把最终指标发给 Tuner，供其更新代理模型、指导下一轮采样。改造方式同样简单，在model.evaluate得到准确率后调用nni.report_final_result(acc)：

def train(args, params): x_train, y_train, x_test, y_test = load_mnist_data(args) model = create_mnist_model(params) model.fit(x_train, y_train, batch_size=args.batch_size, epochs=args.epochs, verbose=1, validation_data=(x_test, y_test), callbacks=[SendMetrics()]) _, acc = model.evaluate(x_test, y_test, verbose=0) nni.report_final_result(acc)

关于两个上报 API 的取值约定，nni/trial.py 的源码 docstring 给出了权威说明：

• metric可以是float，也可以是包含default键（值为 float）的字典；
• 若传字典，Tuner 使用metric['default']，其余字段可在 Web 门户中可视化；
• report_intermediate_result内部以type='PERIODICAL'发送指标，并为每次上报递增序列号；report_final_result则以type='FINAL'、sequence=0发送；
• 两个 API 都断言了nni.get_next_parameter()必须在此之前被调用过（否则在 NNI 平台上会直接断言失败），这是 Trial 代码必须遵循的调用顺序。

此外，report_intermediate_result与report_final_result支持同时传更多自定义字段用于 WebUI 展示；也可以调用 nni/trial.py 中的get_trial_id()、get_sequence_id()等在日志或上报中标记当前 Trial 身份。

完整示例：从零到可运行的 NNI Trial

将以上四步合并，即得到一份完整、可直接运行的 NNI Trial（对应 examples/trials/mnist-keras/mnist-keras.py）：

import argparse import logging import keras import numpy as np from keras import backend as K from keras.datasets import mnist from keras.layers import Conv2D, Dense, Flatten, MaxPooling2D from keras.models import Sequential import nni LOG = logging.getLogger('mnist_keras') K.set_image_data_format('channels_last') H, W = 28, 28 NUM_CLASSES = 10 def create_mnist_model(hyper_params, input_shape=(H, W, 1), num_classes=NUM_CLASSES): layers = [ Conv2D(32, kernel_size=(3, 3), activation='relu', input_shape=input_shape), Conv2D(64, (3, 3), activation='relu'), MaxPooling2D(pool_size=(2, 2)), Flatten(), Dense(100, activation='relu'), Dense(num_classes, activation='softmax') ] model = Sequential(layers) if hyper_params['optimizer'] == 'Adam': optimizer = keras.optimizers.Adam(lr=hyper_params['learning_rate']) else: optimizer = keras.optimizers.SGD(lr=hyper_params['learning_rate'], momentum=0.9) model.compile(loss=keras.losses.categorical_crossentropy, optimizer=optimizer, metrics=['accuracy']) return model def load_mnist_data(args): (x_train, y_train), (x_test, y_test) = mnist.load_data() x_train = (np.expand_dims(x_train, -1).astype(float) / 255.)[:args.num_train] x_test = (np.expand_dims(x_test, -1).astype(float) / 255.)[:args.num_test] y_train = keras.utils.to_categorical(y_train, NUM_CLASSES)[:args.num_train] y_test = keras.utils.to_categorical(y_test, NUM_CLASSES)[:args.num_test] return x_train, y_train, x_test, y_test class SendMetrics(keras.callbacks.Callback): def on_epoch_end(self, epoch, logs={}): LOG.debug(logs) nni.report_intermediate_result(logs) def train(args, params): x_train, y_train, x_test, y_test = load_mnist_data(args) model = create_mnist_model(params) model.fit(x_train, y_train, batch_size=args.batch_size, epochs=args.epochs, verbose=1, validation_data=(x_test, y_test), callbacks=[SendMetrics()]) _, acc = model.evaluate(x_test, y_test, verbose=0) LOG.debug('Final result is: %d', acc) nni.report_final_result(acc) def generate_default_params(): return { 'optimizer': 'Adam', 'learning_rate': 0.001 } if __name__ == '__main__': PARSER = argparse.ArgumentParser() PARSER.add_argument("--batch_size", type=int, default=200, help="batch size", required=False) PARSER.add_argument("--epochs", type=int, default=10, help="Train epochs", required=False) PARSER.add_argument("--num_train", type=int, default=1000, help="Number of train samples to be used, maximum 60000", required=False) PARSER.add_argument("--num_test", type=int, default=1000, help="Number of test samples to be used, maximum 10000", required=False) ARGS, UNKNOWN = PARSER.parse_known_args() try: RECEIVED_PARAMS = nni.get_next_parameter() LOG.debug(RECEIVED_PARAMS) PARAMS = generate_default_params() PARAMS.update(RECEIVED_PARAMS) train(ARGS, PARAMS) except Exception as e: LOG.exception(e) raise

配套的搜索空间与实验配置

要让上述 Trial 真正在 NNI 中参与超参数搜索，还需要两个配套文件：搜索空间描述与实验配置。

搜索空间 examples/trials/mnist-keras/search_space.json 定义了 Tuner 可在哪些参数上采样：

{ "optimizer":{"_type":"choice","_value":["Adam", "SGD"]}, "learning_rate":{"_type":"choice","_value":[0.0001, 0.001, 0.002, 0.005, 0.01]} }

optimizer在Adam与SGD之间选择，learning_rate在 5 个离散值中选择，与 Trial 代码中create_mnist_model(hyper_params)读取的键一一对应——搜索空间的 key 必须与hyper_params中的字段名完全一致，Tuner 采样结果才能被PARAMS.update(RECEIVED_PARAMS)正确覆盖。

实验配置 examples/trials/mnist-keras/config.yml 则声明了实验的运行方式：

authorName: default experimentName: example_mnist-keras trialConcurrency: 1 maxExecDuration: 1h maxTrialNum: 10 #choice: local, remote, pai trainingServicePlatform: local searchSpacePath: search_space.json #choice: true, false useAnnotation: false tuner: #choice: TPE, Random, Anneal, Evolution, BatchTuner, MetisTuner #SMAC (SMAC should be installed through nnictl) builtinTunerName: TPE classArgs: #choice: maximize, minimize optimize_mode: maximize trial: command: python3 mnist-keras.py codeDir: . gpuNum: 0

关键字段含义如下：

• trainingServicePlatform: local：在本地运行 Trial（可选值包括 local、remote、pai，以及仓库中对应的 config_pai.yml 等变体）；
• searchSpacePath：指向搜索空间文件；
• tuner.builtinTunerName: TPE：使用内置 TPE 算法（可选 TPE、Random、Anneal、Evolution、BatchTuner、MetisTuner 等，SMAC 需通过 nnictl 另行安装）；
• tuner.classArgs.optimize_mode: maximize：Tuner 按最大化方向优化，因为 Trial 上报的是 accuracy；若上报的是 loss 则应设为minimize；
• trial.command与trial.codeDir：声明如何启动 Trial 及其代码目录，因此 Trial 的入口脚本名、相对路径必须与之一致；
• trial.gpuNum：每个 Trial 分配的 GPU 数量，本地调试可设为 0。

配置就绪后，即可用nnictl create --config config.yml创建实验，NNI 会启动 nni manager 并按maxTrialNum/maxExecDuration调度多个 Trial 并发执行。

从源码理解数据流：Trial 与 NNI 框架如何通信

以上 API 的背后是一条清晰的通信链路。从源码结构看，Trial 侧所有上报与接收操作最终都收敛到命令通道（Command Channel）抽象：

• 抽象基类 nni/runtime/trial_command_channel/base.py 定义了receive_parameter()（从 NNI manager 接收参数记录）与send_metric()（发送指标，类型限定为PERIODICAL或FINAL，最终指标序号必须为 0）两个接口；
• 不同运行环境下有不同实现，如 standalone.py（脱离 NNI 运行时使用）、local_legacy.py（本地传统模式）、v3.py（NNI v3 新式通道）等；
• nni/trial.py 中的get_next_parameter、report_intermediate_result、report_final_result均为对这些通道实现的薄封装，并通过trial_env_vars（定义于 nni/runtime/env_vars.py，包括NNI_EXP_ID、NNI_TRIAL_JOB_ID、NNI_TRIAL_SEQ_ID、NNI_OUTPUT_DIR等）获得当前实验与 Trial 的上下文。

可以推断，整个数据流为：Tuner 依据搜索空间采样 → 参数记录经命令通道下发至 Trial → Trial 调用get_next_parameter()获取 → 训练过程中按周期调用report_intermediate_result()供 Assessor 决策 → 训练结束调用report_final_result()将最终指标传回 Tuner 更新其模型，形成闭环。这一闭环正是 NNI 超参搜索能够"越搜越好"的底层机制。

小结与编写规范速查

编写一个运行在 NNI 上的 Trial 可归纳为四个步骤：准备可运行的原始代码 → 用nni.get_next_parameter()获取 Tuner 参数并合并进默认参数 → 用nni.report_intermediate_result()周期上报中间结果 → 用nni.report_final_result()上报最终结果。在此之上，遵循以下规范可以显著减少踩坑：

1. 每个 Trial 只调用一次get_next_parameter()；
2. 上报指标优先使用 float，或包含 float 类型default键的字典；
3. report_intermediate_result/report_final_result必须在get_next_parameter()之后调用；
4. 搜索空间的字段名与 Trial 代码读取的超参键保持一致；
5. 用try/except包裹主逻辑并记录异常，便于在 NNI Web 门户中定位失败原因；
6. 利用NNI_OUTPUT_DIR等环境变量管理输出文件（日志、模型、TensorBoard 事件），保证多 Trial 并发互不干扰。

更完整的示例（PyTorch、TensorFlow 等不同框架）可继续阅读 examples/trials/README.md 同目录下的 examples/trials/ 其他示例，以及 docs/source/hpo/quickstart.rst 等官方文档。

【免费下载链接】nniAn open source AutoML toolkit for automate machine learning lifecycle, including feature engineering, neural architecture search, model compression and hyper-parameter tuning.项目地址: https://gitcode.com/gh_mirrors/nn/nni