本章涵盖
- 使用 Kubernetes PersistentVolumes 存储与检索数据集
- 使用 MLflow 与 TensorBoard 追踪并可视化训练过程
- 理解血缘（lineage）与实验追踪的重要性
在生产级 ML 系统中，高效的模型训练不仅仅取决于算法与数据集——它还需要用于数据管理、实验追踪与模型版本管理的健壮基础设施（图 9.1）。第 8 章侧重于构建基础训练流水线，而本章将直面把这些流水线扩展到生产使用时所遇到的挑战。通过身份证检测与电影推荐系统的动手示例，我们将探索如何借助 Kubernetes PersistentVolumes（PVs）高效管理大规模数据集，如何用 MLflow 与 TensorBoard 系统化追踪实验，以及如何维护清晰的模型血缘以支持生产部署。
图 9.1 心智地图：我们继续聚焦于流水线的第 4 与第 5 个步骤/组件——模型训练（4）与评估（5）
随着我们构建的模型越来越多、模型生命周期中的干系人数量不断增加，训练过程中的可追溯性与异步可观测性就变得更加重要。我们还会深入 TensorBoard，它能提升对模型训练环节的可见性；随后把重点切换到 MLflow，它支持血缘与模型版本管理。虽然训练模型并不严格依赖这些能力，但本章会深入讲解一些概念，帮助我们更从容地把模型反复且确定性地交付到生产环境。
我们先从讨论一种在流水线中访问数据的不同方式开始：Kubeflow Pipelines 的 Kubernetes SDK。
9.1 使用 PersistentVolumeClaim 存储数据
我们之前的数据处理方式虽然能用，但存在几个限制：
- 各组件之间会重复下载同一份数据集
- 存储利用率低下
- 流水线步骤之间的数据共享能力有限
Kubernetes PersistentVolumeClaims（PVCs）为生产工作流提供了更稳健的解决方案。通过为流水线组件提供一个持久且可共享的存储空间，PVC 能帮助我们做到：
- 数据只下载一次并被高效复用
- 组件之间无缝共享数据
- 更有效地管理存储资源
到目前为止，我们使用流式读取（stream reads）和直接连接 MinIO 的方式来访问数据集。继续之前，我们需要先安装一个库：
% pip install kfp[kubernetes]
我们将用这个库在 Kubeflow 流水线中访问 Kubernetes 相关特性。本例中，我们会创建一个 PersistentVolumeClaims 组件，用于动态创建 PVC、挂载 PVC，最后再删除 PVC。
9.1.1 使用 PVC 重构流水线
我们在第 7、8 章中处理流水线数据的方式并不理想，具体来说，当前做法是这样的：
- 在 Download dataset 组件中下载完整数据集。
- 在 Train model 组件中重新下载训练/验证/测试切分后的数据集。
- 在 Validate model 组件中再次重新下载验证切分的数据集。
我们在第 4–6 章已经讨论过一种更好的做法，但这里我们快速回顾一下这些概念，以便与项目流水线对齐。
9.1.2 更高效的数据集管理
处理大规模数据集时，必须高效管理。理想情况下，数据集只需要下载一次，然后通过 PV 来访问。这样训练与测试时数据随时可用，不需要反复下载。
但 MinIO 不是传统的可移植操作系统接口（POSIX）文件系统，因此很难像普通文件系统那样轻松复制或移动文件。为了解决这个问题，我们引入 PVC，使我们能够以更灵活的方式操作数据集（图 9.2）。通过使用 PVC，我们可以创建依赖同一个 PVC 的组件，让数据集管理更高效、更可扩展，确保训练过程中可以无缝访问数据。
图 9.2 所有组件都依赖名为 createpvc 的 PVC。
接下来的几个小节，我们将一步步带你把这个流水线重构为使用 PVC。过程中你会看到代码如何逐渐被简化。我们先从创建一个 volume operator（在 KFP 的术语里叫 VolumeOp）开始，它会为流水线创建 PVC。
9.1.3 创建 VolumeOp
创建 VolumeOp 类似于创建一个组件：先从流水线定义开始，然后创建 VolumeOp。你需要指定名称、资源名称以及模式（mode）。大多数 storage class 支持 ReadWriteOnce，所以这是一个安全的选择。最后，你还可以指定卷大小，如下所示。
清单 9.1 用于创建流水线 PVC 的 VolumeOp
@dsl.pipeline(name="YOLO Object Detection Pipeline")
def pipeline(
epochs: int = 1,
batch: int = 8,
random_state: int = 42,
yolo_model_name: str = "yolov8n_custom"
):
pvc_name = 'yolo-pipeline-pvc'
pvc = kubernetes.CreatePVC( #1
pvc_name=pvc_name,
access_modes=['ReadWriteOnce'], #2
size='1Gi', #3
storage_class_name='local-path', #4
)
#1 创建 pvc 组件，底层会创建一个 PVC
#2 指定卷的访问模式为 ReadWriteOnce
#3 指定卷大小
#4 设置 storage class（注意这取决于集群配置！）
接下来我们看下游组件如何使用 PVC，从 download_op 开始。
9.1.4 在 Download Op 中使用 PVC
要使用 PVC，你需要调用 kfp.kubernetes.mount_pvc()。该函数接收前面创建的 pvc 组件名称，以及挂载点（mount point）。当我们往卷里写数据时，就会引用这个挂载点，如下所示。
清单 9.2 在 download_op 中挂载 PVC
from kfp import kubernetes #1
def pipeline(...):
pvc = ...
download_op = download_task()
kubernetes.mount_pvc( #2
download_op, #2
pvc_name=pvc_name, #2
mount_path='/data' #2
) #2
#1 导入 mount_pvc 函数
#2 使用 PVC 名称与挂载点挂载 PVC
更有意思的是：用于创建组件的 download_dataset()（函数定义）会怎么变化。我们已经可以大幅简化这个函数了。现在不再需要 MinIO 相关代码，因为我们只需把数据集下载到 pvc 创建的卷里即可。记住：这只是处理数据的另一种方式，并不意味着 MinIO 组件一定过时或不能用。选择哪种方法取决于问题定义、模型类型以及基础设施限制。下面的清单展示了把数据集下载并解压到卷挂载路径的做法。
清单 9.3 使用卷挂载点下载数据集
@dsl.component(
packages_to_install=["requests", "boto3", "tqdm"],
base_image="python:3.11"
)
def download_dataset():
import requests
import tarfile
from tqdm import tqdm
base_url = "https://manning.box.com/shared/static"
url = f"{base_url}/coiv3n2t5t0v42xgfhlsfi8bhvd7b441.gz"
downloaded_file = "DATASET.gz"
response = requests.get(url, stream=True)
file_size = int(response.headers.get("Content-Length", 0))
progress_bar = tqdm(total=file_size, unit="B", unit_scale=True)
with open(downloaded_file, 'wb') as file:
for chunk in response.iter_content(chunk_size=1024):
progress_bar.update(len(chunk))
file.write(chunk)
with tarfile.open(downloaded_file, 'r:gz') as tar:
tar.extractall("/data/DATASET") #1
#1 将文件解压到 /data/DATASET 下的目录
注意：现在我们甚至不需要显式创建目录了！接下来我们看下有了 pvc 之后如何进行数据集切分。
9.1.5 直接切分数据集
为 split_dataset_op 启用卷挂载与刚才相同。下面清单展示了具体做法。
清单 9.4 在 split_dataset_op 中挂载卷
def pipeline(...):
download_op = ...
split_op = split_dataset(
random_state=random_state,
).after(download_op) #1
kubernetes.mount_pvc( #2
split_op, #2
pvc_name=pvc_name, #2
mount_path='/data' #2
) #2
#1 显式指定 split_op 在 download_op 之后运行
#2 使用 PVC 名称与挂载点挂载 PVC
我们对 split_op 的挂载方式完全一样。这里的主要差异在于：我们需要调用 .after(download_op) 来显式定义执行顺序。此前章节使用 Artifacts 时，存在显式的数据依赖，因此 Kubeflow 能推断组件顺序；但采用这种方式后，需要用 .after() 把顺序写明。
split_dataset_op 的实现代码也变得更简单了。我们不再为不同 split 写出图像与标签文件名，而是更直接地把文件移动到卷内的不同目录下。
清单 9.5 使用 VolumeOp 简化后的 split_dataset 函数
@dsl.component(
packages_to_install=["scikit-learn"],
base_image="python:3.11"
)
def split_dataset(random_state: int):
import os
import glob
import shutil
from sklearn.model_selection import train_test_split
BASE_PATH = "MINIDATA"
images = list(
glob.glob(
os.path.join(
"/data/DATASET",
BASE_PATH,
"images", "**"
))) #1
labels = list(
glob.glob(
os.path.join(
"/data/DATASET",
BASE_PATH,
"labels", "**"
))) #2
train_ratio = 0.75
validation_ratio = 0.15
test_ratio = 0.10
x_train, x_test, y_train, y_test = train_test_split(
images,
labels,
test_size=1 - train_ratio,
random_state=random_state
)
x_val, x_test, y_val, y_test = train_test_split(
x_test, y_test,
test_size=test_ratio / (test_ratio + validation_ratio),
random_state=random_state
)
for split in ["train", "test", "val"]:
for category in ["images", "labels"]:
os.makedirs(
os.path.join(
"/data", split, category),
exist_ok=True) #3
def move_files(
files,
split,
category): #4
for src in files:
dest = os.path.join(
"/data",
split,
category,
os.path.basename(src)) #4
shutil.copy2(src, dest) #4
move_files(x_train, "train", "images") #5
move_files(y_train, "train", "labels") #5
#1 列出 /data/DATASET 下的所有图像
#2 列出 /data/DATASET 下的所有标签
#3 在 /data/DATASET 下创建各 split 的目录
#4 把文件移动到对应 split（而不是复制）
#5 调用 move_files，把文件写入各 split；test 与 val 同理
9.1.6 简化模型训练
这并不令人意外：通过使用 mount_pvc 在 train_model_op 中挂载卷，我们显著简化了模型训练流程，如下一个清单所示。使用 PVC 管理数据集后，我们不再需要从 MinIO 手动下载文件，也不用纠结数据配置 YAML 的路径问题。
清单 9.6 在 train_model_op 中挂载卷
def pipeline(...):
train_op = train_model(
epochs=epochs,
batch=batch,
yolo_model_name=yolo_model_name,
).after(split_op)
kubernetes.mount_pvc( #1
op,
pvc_name=pvc_name,
mount_path='/data'
)
#1 使用 PVC 名称与挂载点挂载 pvc 组件
与此同时，train_model 的函数也被大幅简化。特别是：函数签名不再接收 6 个 InputPath 参数，且 MinIO 下载文件的代码也消失了。注意：我们在创建数据配置 YAML 文件时也更新了路径。剩下的主要就是创建 YAML 并训练模型的代码（下一个长清单展示了完整的简化版 train_model）。
清单 9.7 使用卷的简化版 train_model
@dsl.container_component
def train_model( #1
epochs: int,
batch: int,
yolo_model_name: str,
model_output: dsl.Output[dsl.Model],
data_yaml: dsl.Output[dsl.Artifact]
):
return dsl.ContainerSpec(
image='python:3.11-slim',
command=['bash', '-c'],
args=[
f'''
# #2
...
cat << 'EOF' > /train.py
{TRAINING_SCRIPT}
EOF
python3 /train.py \ #3
--epochs "$0" \
--batch "$1" \
--model-name "$2" \
--model-output "$3" \
--data-yaml "$4"
''',
epochs,
batch,
yolo_model_name,
model_output.path,
data_yaml.path
]
)
#1 移除了训练/测试/验证数据集的输入参数
#2 安装系统依赖
#3 从训练脚本中移除了数据集相关参数
同时别忘了在 TRAINING_SCRIPT 中移除对数据集路径的引用：
TRAINING_SCRIPT = '''
def parse_args(): #1
parser = argparse.ArgumentParser(description='Train YOLO model')
parser.add_argument(
'--epochs',
type=int,
required=True,
help='Number of epochs')
parser.add_argument(
'--batch',
type=int,
required=True,
help='Batch size')
parser.add_argument(
'--model-name',
required=True,
help='Name of the model')
parser.add_argument(
'--model-output',
required=True,
help='Path to save the model')
parser.add_argument(
'--data-yaml',
required=True,
help='Path to save data.yaml')
return parser.parse_args()
def main():
args = parse_args()
data = {
'train': '/data/train/images', #2
'val': '/data/val/images', #2
'test': '/data/test/images', #2
'nc': 1,
'names': {
0: 'id_card'
}
}
with open(data_yaml_path, 'w') as file:
yaml.dump(data, file)
model = YOLO('yolov8n.pt')
results = model.train(...)
if __name__ == "__main__":
main()
'''
#1 移除了训练/测试/验证数据集路径的输入参数
#2 显式引用 PVC 中的数据集路径
除了移除参数之外，其他都保持不变。
9.1.7 简化模型验证
模型验证也是同样的思路。如下清单所示，我们从流水线代码开始。
清单 9.8 在 validate_model_op 中挂载卷
def pipeline(...):
validate_op = validate_model(
data_yaml=train_op.outputs["data_yaml"], #1
model=train_op.outputs["model_output"], #2
).after(train_op)
kubernetes.mount_pvc( #3
validate_op, #3
pvc_name=pvc_name, #3
mount_path='/data' #3
) #3
#1 传入前一个组件输出的 data_yaml
#2 传入前一个组件输出的 model_output
#3 使用 PVC 名称与挂载点挂载 pvc 组件
MinIO 相关代码也被移除了。weights_path 使用挂载点作为根路径构造，然后用于初始化模型，如下所示。
清单 9.9 使用卷的简化版 validate_model
@dsl.component(
base_image="ultralytics/ultralytics:8.3.56-cpu",
packages_to_install=["minio", "tqdm"],
)
def validate_model(
data_yaml: Input[Artifact],
model: Input[Model],
metrics: Output[Metrics]
):
from ultralytics import YOLO
import os
model = YOLO(model_path) #1
validation_results = model.val(
data=os.path.join(data_yaml.path, "data.yaml"), #2
imgsz=640,
batch=1,
verbose=True
)
metrics.log_metric("map50-95", validation_results.box.map)
metrics.log_metric("map50", validation_results.box.map50)
metrics.log_metric("map75", validation_results.box.map75)
#1 从输入参数加载模型
#2 YAML 文件引用了 PVC 中定义的路径
正如你看到的，使用 PVC 让我们大幅删除了大量代码，也让数据存储变得简单得多！图 9.3 展示了把整个流水线加载到 Kubeflow Pipelines（KFP）后呈现的样子。
图 9.3 使用 PVC 的完整训练-验证流水线
9.2 使用 TensorBoard 追踪训练过程
在运行生产训练流水线时，对训练过程的可见性至关重要。MLflow 擅长在更高层面追踪实验，而 TensorBoard 则能对训练过程本身提供更细粒度的洞察。我们通常采用的一个好方法是：用 MLflow 追踪高层指标并管理实验；同时用 TensorBoard 做单次 run 的细节追踪与模型调试。我们会先深入 TensorBoard，并更系统地探索它，再拉远视角介绍 MLflow。这样既能突出两者各自的用例与优势，又能让概念推进保持线性。同时，TensorBoard 还原生集成在 Kubeflow 中，这也很方便。
TensorBoard 在构建生产 ML 工作流时非常有用，因为它提供了：
- 对训练深度指标的实时监控
- 对模型行为的可视化调试
- 多次训练运行的对比
- 模型架构可视化
对于耗时较长的训练，你通常希望能够跟踪模型训练的进度。你可能已经熟悉 TensorBoard 这个可视化工具，它可以用来呈现训练的多个方面——从 loss、accuracy 等曲线图，到模型架构本身。
YOLOv8 虽然不是用 TensorFlow 编写的，但同样支持 TensorBoard。Kubeflow 也内置了对 TensorBoard 的支持。不过，要让它在我们的目标检测流水线里工作起来，还需要做一点额外工作。
目前，YOLOv8 库会在 project 目录下生成 TensorBoard 能读取的日志。我们不会再使用 OutputPath（让 Kubeflow 隐式创建路径）的方式，而是传入一个 VolumeOp，定义挂载点，然后把这个挂载点作为 project 路径。
9.2.1 启动一个新的 TensorBoard
在启动 TensorBoard 之前，确保你已经创建了一次流水线运行，并且至少 create-pvc 这个操作已经完成，如图 9.4 所示。
图 9.4 在启动 TensorBoard 之前，确保 VolumeOp 已完成。
因为我们要从已有的 run 挂载 PVC，所以在启动 TensorBoard 时要记下要使用的 PVC 名称。做法是：点击该组件，并在屏幕上的 pvc_name 行记录下这个名字（图 9.5）。记住，在生产环境中，这个名字就是训练组件保存日志的那个 PVC 的名称。
图 9.5 点击组件后会显示流水线的 PVC 名称。
拿到名字后，在 Kubeflow 侧边栏选择 TensorBoards，然后点击 New TensorBoard（图 9.6）。
图 9.6 在 Kubeflow 中创建一个新的 TensorBoard
点击 New TensorBoard 会弹出如图 9.7 所示的对话框。到这里后，按以下步骤操作：
图 9.7 创建新的 TensorBoard 并指向已有的流水线 PVC
- 为 TensorBoard 填写一个名称。这里我们使用
yolo-object-detection。 - 选择 PVC 单选项。
- 选择 PVC Name。PVC 名称的前缀会与任一组件中显示的前缀匹配。
- 将 Mount Path 输入框留空。
- 完成后点击 CREATE 按钮。
过一会儿，Kubeflow 会启动一个 TensorBoard 实例。当你看到绿色对勾，并且 CONNECT 按钮不再是灰色不可点时，就可以点击 CONNECT（图 9.8）。
图 9.8 TensorBoard 成功启动
9.2.2 探索 YOLOv8 的默认图表
连接到 TensorBoard 后，你可以探索 YOLOv8 默认提供的各种图表，例如图 9.9 中展示的那些。这里我们重点展示一段 YOLO 训练组件样例运行生成的图表。这些指标因此是针对该模型与该训练类型的。如果你换用其他模型和/或指标，图表内容可能不同，但可视化的基本形态是相同的。
图 9.9 YOLOv8 库默认提供的一组图表
这些平均精度均值（mean Average Precision，mAP）图表可以帮助你更好理解目标检测模型的表现。通过追踪这些指标，你可以做到：
- 监控模型表现：观察模型的准确率与精确率如何随时间与训练 epoch 变化。这也有助于识别拟合问题以及潜在的早停触发点。
- 识别趋势与模式：分析曲线，寻找超参数、数据质量或其他因素与模型表现之间的相关性。
- 优化模型训练：利用这些洞察调整超参数、尝试不同架构，或对模型进行微调以获得更好的结果。
通过在 TensorBoard 中可视化这些指标，你将更深入地理解目标检测流水线的表现，并基于数据做出改进决策。
除了指标追踪，TensorBoard 还允许我们可视化模型架构。该图表清晰展示了数据如何在 YOLOv8 目标检测模型中流动。
理解模型复杂度
通过可视化模型架构，我们可以更深入地理解它如何处理输入并生成预测。这在处理像 YOLOv8 这样包含多层结构与连接的复杂模型时尤其有用（图 9.10）。
图 9.10 在 TensorBoard 中查看模型架构
9.3 电影推荐项目
在用目标检测系统把基础设施搭建好之后，我们把这些概念应用到另一类 ML 问题——推荐系统。这种切换能展示我们的流水线模式如何适配不同的 ML 任务，同时也会引入一些新的考量点，包括：
- 高效处理结构化数据
- 管理不同类型的模型产物（artifacts）
- 实现合适的验证策略
- 搭建面向生产的实验追踪能力
我们不会把重点放在推荐算法的细节上，而是专注于构建一个支持持续模型改进的稳健训练基础设施。
在上一章里，我们已经开始讨论电影推荐项目的数据与数据管理：我们创建了一个数据流水线，把数据保存为 Parquet 文件。我们不会专门深入数据探索（data exploration），但在正常工作流中，在走到这里之前，通常会先在实验环境（例如 notebook）里做大量分析与建模。由于本书主要讨论的是 ML 的脚手架（scaffolding），我们选择跳过这部分。
截至目前，我们已经收集了一些数据并把它保存到集群上的 MinIO bucket 里，然后把数据切分为 train、test 和 validation。最后，我们还对这些数据做了一些基础的质量分析（rudimentary quality analysis），现在已经准备好开始使用这些数据了。在常见场景里，这通常意味着从一个或多个上游数据源查询数据、做清洗与处理变换，然后把数据保存到数据湖中。你现在掌握的这些基础能力足以实现这一整套流程。
在本节中，我们会在此基础上重点聚焦训练模型。我们将构建组件来检查进入训练流程的数据质量，用这些数据训练模型，定义用于验证模型的指标，最后在 notebook 中试用训练好的模型，为“电影之夜”生成一些推荐。我们还会看看如何使用 MLflow 记录实验并追踪模型版本。虽然你也可以像目标检测例子那样用 Kubeflow 来做这些事，但 MLflow 提供了另一种选择，也是团队工具箱里很值得拥有的工具。
注意 在使用数据训练或评估模型之前，最好再一次验证你对数据的假设。尽管在数据创建流程中已经做过验证，但在训练/评估之前重新验证同样的假设，可以确保数据摄取与变换流水线能够与训练与评估流水线解耦。
9.3.1 从 MinIO 读取数据并做质量保障
我们从编写一个组件开始，用于检查 MinIO bucket 里的数据质量。目前我们会复用上一章的组件，但会根据训练需求增加一些检查项。具体来说，我们会检查是否满足正确的 schema，并验证数据类型。虽然这看起来很简单，但这些检查能捕获大多数可能渗入生产环境的问题，并为我们提供一套可扩展的检查仓库——随着我们发现更多数据缺陷，还可以不断把新的检查项加进去。和 ML 运维中的其他事情一样，识别质量问题也是一个迭代过程。
与前一个例子不同，这里我们尝试直接从 MinIO 读取数据。这种方式能更快开始训练，因为训练 run 一开始不需要做大规模拷贝；同时也让我们不必管理共享的 VolumeOp。再强调一次：无论是训练开始时把整份数据拷贝到一个公共共享卷里，还是直接从 MinIO bucket 读取，两种方式都合理，具体选择很大程度取决于所使用的数据类型。
数据访问的考量（DATA ACCESS CONSIDERATIONS）
把数据下载到本地磁盘的问题在于：它会消耗更多存储空间，并且对于大数据集来说，下载耗时更长，从而导致训练 run 启动延迟更久、存储成本更高。虽然我们会在训练结束后删除数据，但另一种做法是直接从 MinIO bucket 读取数据。这里的权衡在于：这种方式对网络更“吃”，因为所有数据访问都发生在网络上，而不是本地卷。
如果你的现代化基础设施里，各节点之间都有高带宽网络连接，那么这种方式可以在不影响总体运行时间的前提下，加快训练的初始启动时间。务必结合你的基础设施特性认真评估数据访问方式：本地读取存储成本高、启动慢，但训练时更快；网络读取存储成本低、启动快，但要求网络足够快且组件实现需要针对网络访问做优化。
一开始，我们运行数据质量评估（QA），如下所示。先跑一些简单 QA 脚本就足够；后续我们可以把这个组件替换成更复杂的套件，例如 Evidently 或其他数据质量平台。
清单 9.10 数据质量检查
def check_dataset(bucket:str = 'datasets', dataset:str = 'ml-25m'):
from pyarrow import fs, parquet
print("Running QA")
minio = fs.S3FileSystem(
endpoint_override='http://minio-service.kubeflow:9000',
access_key='minio',
secret_key='minio123',
scheme='http')
train_parquet = minio.open_input_file(
f'{bucket}/{dataset}/train.parquet.gzip')
df = parquet.read_table(train_parquet).to_pandas()
is_subset = set(['user_id', 'item_id', 'rating']).issubset(df.columns)
assert is_subset, f'Required columns not found. Found {df.columns}'
print('QA passed!')
提示 在
ml-pipeline-persistenceagentpod 中设置环境变量TTL_SECONDS_AFTER_WORKFLOW_FINISH，可以在定义的 TTL（time to live）后自动清理旧的 VolumeOps。
9.3.2 模型训练组件
对推荐模型，我们选择一个用 PyTorch 编写的简单矩阵分解（matrix factorization）模型（清单 9.11）。这个模型当然谈不上“拿奖”，但它足够简单，可以快速训练出一个可用模型；同时采用 PyTorch 这样的更强框架，也会让我们后续更容易处理模型。
清单 9.11 模型训练代码
def train_model(
model_embedding_factors=20,
model_learning_rate=1e-3,
model_hidden_dims=256,
model_dropout_rate=0.2,
optimizer_step_size=10,
optimizer_gamma=0.1,
training_epochs=30,
train_batch_size=64,
test_batch_size=64,
shuffle_training_data=True,
shuffle_testing_data=True):
import torch
from torch.autograd import Variable
from torch.utils.data import DataLoader
class MatrixFactorization(torch.nn.Module): #1
def __init__(
self,
n_users,
n_items,
n_factors,
hidden_dim,
dropout_rate):
super().__init__()
self.n_items = n_items
self.user_factors = torch.nn.Embedding(
n_users + 1,
n_factors,
sparse=False)
self.item_factors = torch.nn.Embedding(
n_items + 1,
n_factors,
sparse=False)
self.linear = torch.nn.Linear(
in_features=n_factors,
out_features=hidden_dim
)
self.linear2 = torch.nn.Linear(
in_features=hidden_dim,
out_features=1
)
self.dropout = torch.nn.Dropout(p=dropout_rate)
self.relu = torch.nn.ReLU()
def forward(self, user, item):
user_embedding = self.user_factors(user)
item_embedding = self.item_factors(item)
embeddding_vector = torch.mul(user_embedding, item_embedding)
x = self.relu(self.linear(embeddding_vector))
x = self.dropout(x)
rating = self.linear2(x)
return rating
dataset_map = get_datasets_local(split=['train', 'test'])
#1 模型定义类
这个简单的矩阵分解模型（由上面代码片段所示）为用户与物品分别设置了 embedding，并包含两层隐藏层。我们不会在这里深入它的数学细节，但网上有很多资源可供参考。
model = MatrixFactorization(
dataset_map['n_users'],
dataset_map['n_items'],
n_factors=model_embedding_factors,
hidden_dim=model_hidden_dims,
dropout_rate=model_dropout_rate)
optimizer = torch.optim.SGD(
model.parameters(),
lr=model_learning_rate)
scheduler = torch.optim.lr_scheduler.StepLR(
optimizer,
step_size=optimizer_step_size,
gamma=optimizer_gamma)
loss_func = torch.nn.L1Loss()
train_dataloader = DataLoader(
dataset_map['train'],
batch_size=train_batch_size,
shuffle=shuffle_training_data)
test_dataloader = DataLoader(
dataset_map['test'],
batch_size=test_batch_size,
shuffle=shuffle_testing_data)
for train_iter in range(training_epochs):
print(train_iter)
model.train()
t_loss = 0
t_count = 0
for row, col, rating in train_dataloader:
prediction = model(row, col)
loss = loss_func(prediction, rating.unsqueeze(1))
t_loss += loss
t_count += 1
loss.backward()
optimizer.step()
optimizer.zero_grad()
scheduler.step()
model.eval()
te_loss = 0
te_count = 0
print('Evaluating')
with torch.no_grad():
for row, col,rating in test_dataloader:
prediction = model(row, col)
loss = loss_func(prediction, rating.unsqueeze(1))
te_loss += loss
te_count += 1
print(f"Test loss: {te_loss/te_count}")
print(f"Train loss: {t_loss/t_count}")
训练循环也相当简单：使用 batch 输入，并在每次迭代都跑一次测试循环。要加快训练，你可以每 5 次迭代左右才跑一次测试。
9.3.3 评估指标
现在我们有了训练后的模型，就需要定义一些指标。如第 2 章所说，评估指标通常源自业务目标，理想情况下应当在建模前就定义好，以保证客观性并更好对齐干系人；同时，这也有助于 ML 团队跨职能沟通模型到底好（或差）到什么程度——用少量几个数字就能表达。
在我们的场景里，我们要做的是电影推荐引擎。多数情况下，这意味着需要评估：对某个用户而言，推荐列表有多相关（relevant），以及给某个用户提供了多少“正确”的推荐（见清单 9.12）。我们选择了三个基础的评估标准：precision、recall 和均方根误差（RMSE）：
- RMSE——衡量预测评分误差的指标；完美预测时 RMSE 必须为 0，因此 RMSE 越低，模型给出推荐评分的准确性越好。
- Precision——衡量某个用户的一组推荐有多相关，也就是推荐集合“质量”有多好。统计意义上，precision = 真阳性（true positive）数量 /（真阳性 + 假阳性）数量。我们取模型的 top 50 推荐来计算 precision，期望 precision 越高越好。
- Recall——衡量推荐列表中包含多少相关条目。直观地说，recall 更像是有效推荐的“数量”。统计定义上，recall = 真阳性数量 /（假阴性 + 真阳性）数量。同样地，我们用 top 50 推荐来计算 recall，期望 recall 越高越好。
注意 指标是评估工作流中最重要的部分。理想情况下，我们会有几个从不同维度衡量性能的指标，并验证这些指标确实具备良好的描述性。这里我们使用的是玩具指标作为示例，但在生产工作流中，必须非常谨慎地设计并评估指标，以确保它们与业务需求一致。记住：模型表现最终就是由这几个数字来裁决的，指标选错了，要么会让我们错丢一个其实很好的模型，要么会把一个差模型的价值“虚高”。
清单 9.12 模型评估代码
def validate_model(
recommendation_model,
top_k=50,
threshold=3,
val_batch_size=32):
from collections import defaultdict #1
import torch
from sklearn.metrics import mean_squared_error
def calculate_precision_recall(user_ratings, k, threshold):
user_ratings.sort(key=lambda x: x[0], reverse=True)
n_rel = sum(x >= threshold for _, x in user_ratings)
n_rec_k = sum(x >= threshold for x, _ in user_ratings[:k])
n_rel_and_rec_k = sum(
(x >= threshold) and (y >= threshold)
for y, x in user_ratings[:k]
)
precision = n_rel_and_rec_k / n_rec_k if n_rec_k != 0 else 1
recall = n_rel_and_rec_k / n_rel if n_rel != 0 else 1
return precision, recall
user_ratings_comparison = defaultdict(list)
dataset_map = get_datasets_local(split=["val"])
val_dataloader = DataLoader(
dataset_map["val"],
batch_size=val_batch_size,
shuffle=True
)
y_pred = []
y_true = []
recommendation_model.eval()
with torch.no_grad():
for users, movies, ratings in val_dataloader:
output = recommendation_model(users, movies)
y_pred.append(output.sum().item() / len(users))
y_true.append(ratings.sum().item() / len(users))
for user, pred, true in zip(users, output, ratings):
user_ratings_comparison[user.item()].append(
(pred[0].item(), true.item())
)
user_precisions = dict()
user_based_recalls = dict()
k = top_k
for user_id, user_ratings in user_ratings_comparison.items():
precision, recall = calculate_precision_recall(
user_ratings, k, threshold)
user_precisions[user_id] = precision
user_based_recalls[user_id] = recall
average_precision = sum(prec for prec in user_precisions.values())/len(
user_precisions
)
average_recall = sum(rec for rec in user_based_recalls.values())/len(
user_based_recalls
)
rms = mean_squared_error(y_true, y_pred, squared=False)
print(f"precision_{k}: {average_precision:.4f}")
print(f"recall_{k}: {average_recall:.4f}")
print(f"rms: {rms:.4f}")
#1 代码参考：https://mng.bz/eBlJ
9.3.4 使用 MLflow 进行实验追踪
随着 ML 项目规模扩大，追踪实验会越来越困难。像“哪些超参数给出了最好的结果？”或“上个月那次训练 run 有什么不同？”之类的问题，如果没有系统化追踪，就很难回答。MLflow 通过提供以下能力来解决这些挑战：
- 自动化实验日志（logging）
- 集中的模型产物存储
- 标准化的模型打包
- 清晰的模型血缘追踪
我们来看看如何把这些能力集成进训练流水线。现在我们已经有了训练循环，也定义了用于评判模型好坏的指标，接下来就讨论 ML 的一个核心原则：实验追踪。为了说明追踪的重要性，我们先用默认参数跑一个玩具训练任务，看看在这些指标上的表现。
测试 run 参数如下：
- 6,400 条随机数据样本
- train/test 的 batch size 都为 64
- 训练 30 个 epoch
- 其余参数保持默认
表 9.1 给出了结果。
表 9.1 使用默认参数测试玩具样本
| Precision@50 | Recall@50 | RMSE |
|---|---|---|
| 0.7533 | 0.6316 | 0.2904 |
模型还可以，但如果我们改一些超参数再试试呢？我们尝试修改学习率，看看模型表现如何。测试 run 参数如下：
- 6,400 条随机数据样本
- train/test 的 batch size 都为 64
- 学习率为 1e-2
- 训练 30 个 epoch
- 其余参数保持默认
表 9.2 给出了结果。
表 9.2 修改学习率后的测试结果
| Precision@50 | Recall@50 | RMSE |
|---|---|---|
| 0.7422 | 0.8649 | 0.2836 |
这次 run 的表现比之前略差，因此我们想回退并恢复参数。但问题来了：原始学习率是多少？原始那套超参数组合是什么？优化器是不是同一个？
诚然，这是个很简单的例子，我们完全可以通过回滚代码来回答这些问题。但当我们把规模放大，考虑多个大团队在多个天、多个周、甚至几个月里跑大量实验时，问题就会变得明显：我们如何跨时间与团队成员追踪所有实验？我们可以维护一个巨大的表格，但无法保证每个人的实验都被准确记录，也无法确保复现实验所需的一整套超参数都完整无误。当模型在生产中出现异常行为时，这个问题会更严重：要定位根因，我们必须复现训练方案与参数；或者当数据集演进后，我们还需要知道训练时使用的数据集版本。
这就是实验追踪框架登场的地方。你在前一个项目里已经通过 TensorBoard 初步体验过模型对比与训练监控；现在我们会使用另一个框架——MLflow——来跨时间追踪多次运行。
TensorBoard 更侧重理解训练过程本身，而 MLflow 擅长记录与索引关键元信息，包括模型超参数、训练数据集版本、训练代码版本，以及最终训练好的模型。按我们的经验，用 TensorBoard 做训练分析、用 MLflow 做实验追踪与模型版本管理，往往是最佳组合。虽然 MLflow 也提供一些训练指标与实验对比能力，但在丰富的可视化与生态系统方面，它不如 TensorBoard；反过来，MLflow 提供了很强的实验索引与搜索能力，结合模型注册表（model registry）后，MLflow 更适合在大团队里追踪实验。
上传到 MLflow（UPLOADS TO MLFLOW）
从训练代码向 MLflow 上传产物，比看起来要更复杂一些。在常规运行模式下，训练环境需要能独立访问模型产物的存储位置。比如如果我们用 Amazon S3 作为模型存储层，那么训练环境必须配置好访问 S3 的凭证。
为了解决这一点，我们可以将 MLflow tracking server 运行在“代理访问（proxied access）”模式下：模型存储通过 MLflow tracking server 转发，训练代码只与 tracking server 通信。这对大多数用例来说是最省事的配置，但代价是：单一端点要处理大文件上传，可能成为瓶颈。如果你的训练节点配置得当，从性能角度看，通常更好的方式是让训练代码直接与存储后端通信，只把元数据交给 tracking server。
我们先把 MLflow 集成进训练与评估代码。下面是完整代码，用于展示集成细节并突出新增内容。
清单 9.13 带 MLflow 追踪的模型训练代码
def train_model(
mlflow_experiment_name='recommender',
mlflow_run_id=None,
mlflow_tags={},
hot_reload_model_run_id=None,
model_embedding_factors=20,
model_learning_rate=1e-3,
model_hidden_dims=256,
model_dropout_rate=0.2,
optimizer_step_size=10,
optimizer_gamma=0.1,
training_epochs=30,
train_batch_size=64,
test_batch_size=64,
shuffle_training_data=True,
shuffle_testing_data=True):
import torch
from torch.autograd import Variable
from torch.utils.data import DataLoader
import mlflow #1
from torchinfo import summary #1
from mlflow.models import infer_signature #1
input_params = {}
for k, v in locals().items(): #1
if k == 'input_params':
continue
input_params[k] = v
class MatrixFactorization(torch.nn.Module):
. . .
if hot_reload_model_run_id is not None: #2
model_uri = f"runs:/{hot_reload_model_run_id}/model" #2
model = mlflow.pytorch.load_model(model_uri) #2
else:
model = MatrixFactorization( . . .
. . .
mlflow.set_experiment(mlflow_experiment_name) #1
with mlflow.start_run(run_id=mlflow_run_id): #1
for k,v in input_params.items():
if 'mlflow_' not in k:
mlflow.log_param(k, v) #1
mlflow.log_param("loss_function",
loss_func.__class__.__name__) #1
mlflow.log_param("optimizer", "SGD") #1
mlflow.log_params({'n_user':
dataset_map['n_users'], 'n_items':
dataset_map['n_items']}) #1
for k,v in mlflow_tags.items():
mlflow.set_tag(k, v) #1
with open("model_summary.txt", "w") as f:
f.write(str(summary(model)))
mlflow.log_artifact("model_summary.txt") #3
model_signature = None
for train_iter in range(training_epochs):
print(train_iter)
model.train()
t_loss = 0
t_count = 0
for row, col, rating in train_dataloader:
prediction = model(row, col)
if model_signature is None:
model_signature = infer_signature(
{'user': row.cpu().detach().numpy(),
'movie': col.cpu().detach().numpy()},
prediction.cpu().detach().numpy()) #3
loss = loss_func(prediction, rating.unsqueeze(1))
. . .
mlflow.log_metric(
"avg_training_loss",
f"{(t_loss/t_count):3f}",
step=train_iter) #1
scheduler.step()
model.eval()
. . .
with torch.no_grad():
for row, col,rating in test_dataloader:
prediction = model(row, col)
loss = loss_func(prediction, rating.unsqueeze(1))
te_loss += loss
te_count += 1
mlflow.log_metric(
"avg_testing_loss",
f"{(te_loss/te_count):3f}",
step=train_iter) #1
print(f"Test loss: {te_loss/te_count}")
print(f"Train loss: {t_loss/t_count}")
mlflow.pytorch.log_model(model,
"model", signature=model_signature) #3
#1 为额外日志新增的代码行
#2 从注册表加载模型
#3 将模型与参数记录到注册表
注释标出了我们新增的所有代码行，用来提供一些有用能力。首先，我们记录所有输入超参数以及当前实验中使用的优化器类型等信息；我们还记录一些指标，比如训练 run 的平均 train/test loss。然后我们加入从 MLflow 加载之前训练模型的能力（如清单 9.13），从而支持 warm restart 训练。该清单还展示了如何记录诸如模型摘要（model summary）这样的产物，如何推断模型的输入/输出签名（signature），以及最后如何带签名把模型记录到 MLflow——这样我们就能直接从 MLflow 加载模型，而无需额外文档。
把 MLflow 引入训练代码后，会立刻带来巨大的收益：超参数的可追溯性让我们能更快迭代，而不必担心手工记实验与参数；我们现在可以直接把模型保存到并从中央服务器取回，模型存储基础设施的负担也从开发者身上转移走了；基于中央存储的概念也让“从之前 run 热启动训练”成为可能；最后，模型的详细文档（输入、输出与架构）意味着：即便我们需要回溯分析某个模型，或快速跑一次推理验证假设，开发者也只需要看 MLflow 即可。
到这里，你可能已经能看到自动追踪实验的内在价值，但我们再往前一步，把追踪代码集成到模型验证里（清单 9.14）。这样我们就能把训练与评估工作流分离，使评估组件更可移植。
清单 9.14 带 MLflow 追踪的模型验证代码
def validate_model(
model_run_id,
top_k=50,
threshold=3,
val_batch_size=32):
from collections import defaultdict
import torch
import mlflow.pytorch
import mlflow
from sklearn.metrics import mean_squared_error
model_uri = f"runs:/{model_run_id}/model" #1
recommendation_model = mlflow.pytorch.load_model( #1
model_uri) #1
. . .
precision_sum = sum(prec for prec in user_precisions.values())
average_precision = precision_sum / len(user_precisions)
recall_sum = sum(rec for rec in user_based_recalls.values())
average_recall = recall_sum / len(user_based_recalls)
rms = mean_squared_error(y_true, y_pred, squared=False)
. . .
mlflow.log_metric(
f"precision_{k}",
average_precision,
run_id=model_run_id) #2
mlflow.log_metric(
f"recall_{k}",
average_recall,
run_id=model_run_id) #3
mlflow.log_metric(
"rms",
rms,
run_id=model_run_id) #3
#1 从 MLflow 加载模型
#2 将指标写入实验
#3 将指标写入实验
在上面的代码中，我们直接从 MLflow 加载模型，无需原始的模型类，也不需要手动下载权重。MLflow 在中央服务器上管理所有 runs，使我们可以通过一次 API 调用就引用模型权重与指标。该方法以 run ID 作为输入，从而能直接从 MLflow 拉取模型做推理。run ID 可以直接在 MLflow 中获得，也可以来自训练方法的输出（如清单 9.14 所示）。
我们还把指标写回实验。这里值得注意的一点是：即便实验已经完成，我们仍然可以往同一个实验里继续写指标。这在训练与评估流程分离、或模型在初次训练后又被进一步微调等场景里很有用。保持实验与模型标识一致，有助于维护干净的线性历史。
现在我们已经把 MLflow 集成进代码，接下来试着跑一下。你会看到：一旦训练循环开始，Experiments 标签页下就会创建一个新的 run，包含一堆有趣字段。点击该 run 会显示更详细信息，比如训练参数以及指标随训练迭代推进的变化。等训练完成后，你还会看到被记录的模型，以及一些示例代码来运行模型。查看与该 MLflow run 一起记录的 Model 标签页，就能看到这些内容。
要真正体会 MLflow 的优势，我们再改一些超参数，多跑几个实验。我们复用前面修改学习率到 1e-2 的例子：先用默认参数训练并评估指标；然后改学习率，再重复一次流程。
当你跑完两次实验后，选中多个实验并点击 Compare。此时，超参数变化对模型指标的影响、哪些参数组合效果更好，就都能很容易地看出来了。我们甚至可以给 runs 打标签、写描述，确保大团队能理解实验上下文，并通过 tags 让一组 runs 可检索。这在做大规模实验（例如用 Katib 或其他超参搜索工具）、团队协作跑实验、以及向干系人快速解释模型对比时都非常有帮助。
拥有一个模型性能仪表盘，也让组织获得一个“单一玻璃面板（single pane of glass）”：大家可以直接看到实验结果，而不必等数据科学家或 ML 工程师来汇报模型表现。按我们的经验，这种异步沟通模型状态的方式往往是最好的，因为把链接发出去，本质上就“自带”状态更新。图 9.11 展示了一个比较两个模型指标与性能的仪表盘示例，能让干系人在一个页面里快速分析模型。
图 9.11 MLflow 中模型仪表盘与指标示例
9.3.5 使用 MLflow 的模型注册表（Model registry）
在生产级 ML 系统中，模型在部署前通常会经历多个阶段——从实验版本到 staging，再到 production。模型注册表提供了支撑基础设施，用于完成以下任务：
- 系统化追踪模型版本
- 控制模型在不同阶段之间的晋升（promotion）
- 为审计维护清晰的血缘（lineage）
- 需要时支持快速回滚
MLflow 的模型注册表通过模型管理能力帮助我们落地这些实践。我们已经看到 MLflow 能记录训练好的模型，但我们还可以更进一步：在 MLflow 中管理模型的完整生命周期。在高度自动化的 MLOps 体系里（端到端训练与部署都自动化），这种能力尤其强大；不过即便如此，在模型从开发到生产、再到 staging 的晋升过程中，通常仍需要一个人工检查点（manual checkpoint）。
模型注册表也充当所有模型的“单一事实来源”（single source of truth），并提供用于操作模型的 API，从而让管理以及诸如模型服务（model serving）这类复杂任务更简单。最后，模型注册表把强可追溯性内建进来，使团队在将模型部署到生产时，能对模型血缘更有信心。
首先，我们需要理解 MLflow 如何处理模型。在实验阶段，模型会与其超参数和代码一起，作为一次实验 run 的一部分被记录下来。正常情况下，一个问题往往会对应大量实验：既包括手工试验，也包括诸如超参搜索之类的自动化方法。一旦某个模型满足标准，你就会决定把那些有希望的模型单独“拎出来”。随后，每个模型会以一个唯一别名注册到模型注册表中，例如 dev.ml_components.recommender。被注册的模型拥有唯一名称、版本，以及可选的别名和其他关联元数据。
当某个模型满足生产标准后，它可以被拷贝到另一个模型中，例如 prod.ml_components.recommender。如你所料，这个生产模型也会随着时间拥有多个版本，因此我们会用模型别名指向某个特定版本，以标识它当前处于“正在使用”的状态。这个工作流既支持良好的血缘追踪，也允许我们在模型晋升时设置明确的质量门禁（quality gates），从而基于不同标准推动模型晋升。最后，它还支持通过改变生产别名所指向的版本来实现快速回滚。
注册模型的第一种方式是通过 Web UI。该界面提供了一种方便的方式手动注册已记录的模型。试着用 MLflow UI 创建一个名为 Staging 的模型。关于如何创建模型的细节，请参考第 4 章（清单 4.5）。
现在我们已经创建了一个 staging 模型占位符，接下来编写一个组件，用于把模型注册到 MLflow 中某个命名 stage，如清单 9.15 所示。在这个组件里，我们可以把 stage 名称作为输入参数，从而让组件可复用。
清单 9.15 将模型晋升到 staging
def promote_model_to_staging(
new_model_run_id,
registered_model_name="recommender_production",
rms_threshold=0.0,
precision_threshold=-0.3,
recall_threshold=-0.2,
):
import mlflow.pytorc
import mlflow
from mlflow import MlflowClient
from mlflow.exceptions import RestException
client = MlflowClient()
current_staging = None
try:
current_staging = client.get_model_version_by_alias(
registered_model_name, "staging"
)
except RestException:
print("No staging model found. Upgrade current run to staging.")
if current_staging.run_id == new_model_run_id:
print("Input run is already the current staging.")
return
if current_staging is not None:
current_staging_model_data = client.get_run(
current_staging.run_id
).data.to_dictionary()
staging_metrics = current_staging_model_data["metrics"]
new_model_data = client.get_run(new_model_run_id)
new_metrics = new_model_data.data.to_dictionary()["metrics"]
rms_diff = new_metrics["rms"] - staging_metrics["rms"]
prec_key = "precision_50"
rec_key = "recall_50"
prec_diff = new_metrics[prec_key] - staging_metrics[prec_key]
rec_diff = new_metrics[rec_key] - staging_metrics[rec_key]
if rms_diff > rms_threshold:
return
if prec_diff < precision_threshold:
return
if rec_diff < recall_threshold:
return
result = mlflow.register_model(
f"runs:/{new_model_run_id}/model", "recommender_production"
) #1
client.set_registered_model_alias(
"recommender_production", "staging", result.version
) #1
#1 注册并把模型晋升到 staging
该清单只把新训练的模型与当前 staging 模型做对比。其思路是：在一些场景中，晋升到 production 必须是带评审的人工流程；但如果组织更偏好高度自动化的设置，我们也可以添加一些经验性（empirical）的模型测试，并像清单 9.15 那样直接把模型晋升到 production。production 与 staging 的“标识”在模型注册表中本质上只是别名（aliases），工作方式完全一致。下一章里，我们会在构建推理引擎时使用这些 tag 来检索模型，从而完成从数据到推理的整条链路。
晋升模型的另一种方式是使用 API，它与 UI 的体验非常接近。MLflow 官方文档对该流程有描述。
9.3.6 由组件组装成流水线
和上一章一样，我们现在把这些函数编译成组件并连接起来，形成一条流水线。我们不会对这条流水线展开太多细节，因为我们已经花了不少篇幅讨论各个组件。完整流水线见下方清单。
清单 9.16 完整流水线：编译组件
import kfp.dsl as dsl
import kfp
from data_components import qa_data
from training_and_validation_components import (
negative_sampling, get_dataset_metadata,
get_test_valid_dataset,
promote_model_to_staging,
validate_model,
train_model
)
@dsl.pipeline(
name='Model training pipeline',
description='A pipeline to train recommender models on movielens'
)
def training_pipeline(
minio_bucket:str='datasets',
number_of_negative_samples: int = 10,
training_dataset_name:str = 'ml-25m',
training_batch_size: int = 32, #64
training_learning_rate:float = 0.001,
model_embedding_factors: int = 5, #20
model_hidden_dims:int = 64, #256
training_epochs:int = 30,
optimizer_step_size: float= 10.0,
optimizer_gamma: float = 0.1,
model_dropout_rate:float = 0.2,
testing_batch_size: int = 32, #64
shuffle_training_data:bool =True,
shuffle_testing_data:bool =True,
hot_reload_model_id: str = 'none',
validation_top_k:int = 50,
validation_threshold:int = 3,
validation_batch_size: int = 32,
model_promote_rms_threshold: float = 0.0001,
model_promote_precision_threshold: float = -0.3,
model_promote_recall_threshold:float = -0.2,
mlflow_experiment_name: str = 'recommender',
mlflow_registered_model_name: str = 'recommender_production',
mlflow_uri: str='http://192.168.1.90:8080'):
qa_op = qa_data(bucket=minio_bucket).set_display_name("qa-training-data")
dataset_metadata = get_dataset_metadata(
bucket=minio_bucket,
dataset_name=training_dataset_name).after(qa_op)
negative_sampled_data = negative_sampling(
bucket=minio_bucket,
dataset_name=training_dataset_name,
split='train',
num_ng_test=number_of_negative_samples)\
.after(dataset_metadata)\
.set_caching_options(False)
aux_data = get_test_valid_dataset(
bucket=minio_bucket,
dataset_name=training_dataset_name)\
.after(negative_sampled_data)\
.set_caching_options(False)
training = train_model(
mlflow_experiment_name=mlflow_experiment_name,
mlflow_run_id="",
mlflow_tags={},
hot_reload_model_run_id=hot_reload_model_id,
model_embedding_factors=model_embedding_factors,
model_learning_rate=training_learning_rate,
model_hidden_dims=model_hidden_dims,
model_dropout_rate=model_dropout_rate,
optimizer_step_size=optimizer_step_size,
optimizer_gamma=optimizer_gamma,
training_epochs=training_epochs,
train_batch_size=training_batch_size,
test_batch_size=testing_batch_size,
training_data=negative_sampled_data\
.outputs['negative_sampled_dataset'],
training_data_metadata=dataset_metadata.output,
testing_data=aux_data.outputs['testing_dataset'],
shuffle_training_data=shuffle_training_data,
shuffle_testing_data=shuffle_testing_data,
mlflow_uri=mlflow_uri)\
.after(negative_sampled_data)\
.set_caching_options(False)
val = validate_model(
model_run_id=training.output,
top_k=validation_top_k,
threshold=validation_threshold,
val_batch_size=validation_batch_size,
validation_dataset=aux_data.outputs['validation_dataset'],
mlflow_uri=mlflow_uri).after(training).set_caching_options(False)
promote_model_to_staging(
model_run_id=training.output,
registered_model_name=mlflow_registered_model_name,
top_k=validation_top_k,
rms_threshold=model_promote_rms_threshold,
precision_threshold=model_promote_precision_threshold,
recall_threshold=model_promote_recall_threshold,
mlflow_uri=mlflow_uri).after(val).set_caching_options(False)
if __name__ == "__main__":
kfp.compiler.Compiler().compile(
pipeline_func=training_pipeline,
package_path='compiled_pipelines/training_pipeline.yaml')
清单 9.16 到这里你应该已经很熟悉了：我们把所有函数编译成组件，再把组件组合成流水线。然后编译流水线，并用 MLflow 与 MinIO 管理模型与数据。该流水线输入参数很多，因为它做的事情也不少。你可以把它拆开，用一条流水线去触发另一条流水线。在 Kubeflow 中，流水线本身也可以作为组件再组合起来。
9.3.7 在 notebook 中做本地推理
在把模型部署到生产之前，在一个可控环境中验证它们至关重要。本地推理测试可以帮助我们做到：
- 用特定测试用例验证模型行为
- 调试意料之外的推荐结果
- 向干系人展示模型能力
- 建立对训练流水线的信心
现在我们已经有了训练好的模型，并假设我们已经把其中一个晋升到 production，那么就来跑一跑这个模型！我们会写一个快速脚本，通过 run 引用模型，然后传入一个用户生成推荐，如下所示。
清单 9.17 在 notebook 中推理
import torch
import mlflow
from mlflow import MlflowClient
import pandas as pd
import os
class RecommendationSystem:
def __init__(self, registered_model_name, device='cpu'):
client = MlflowClient()
current_prod = client.get_model_version_by_alias(
registered_model_name, "prod"
) #1
#1
model_uri = f"runs:/{current_prod.run_id}/model" #1
self.model = mlflow.pytorch.load_model(model_uri) #1
self.device = device
self.model.to(self.device)
self.movie_map = self.generate_map('/Users/shanoop/Downloads/ml-25m')
def generate_map (self, path):
names = ['movie_id', 'title', 'genres']
ratings_df = pd.read_csv(
os.path.join(path, 'movies.csv'),
names=names,
index_col=False,
skiprows=1
)
return ratings_df
def movieID_to_name(self, ids):
movies = self.movie_map.loc[self.movie_map['movie_id'].isin(ids)]
return movies
def recommend(self, user_id, top_k=10, ranked_movies=None):
user_id = torch.tensor([user_id], dtype=torch.long).to(self.device)
all_items = torch.arange(
1,
self.model.n_items + 1,
dtype=torch.long
).to(self.device)
if ranked_movies is not None:
ranked_movies = torch.tensor(
ranked_movies,
dtype=torch.long).to(self.device)
unrated_items = all_items[~torch.isin(all_items, ranked_movies)]
else:
unrated_items = all_items
user_ids = user_id.repeat(len(unrated_items))
with torch.no_grad():
self.model.eval()
predictions = self.model(user_ids, unrated_items).squeeze()
top_n_indices = torch.topk(predictions, top_k).indices
recommended_items = unrated_items[top_n_indices].cpu().numpy()
return self.movieID_to_name(recommended_items.tolist())
#1 通过名称从 MLflow 加载模型
清单 9.17 的代码与验证代码非常相似。得益于模型注册表 API，我们不需要模型定义或类。因为我们已经记录了输入与输出 tensor，并用维度对它们进行了命名，最终用户就可以很容易地对模型进行推理。这一点会在下一章体现得更明显：我们将使用 BentoML 和推理框架，直接从 MLflow 中的模型定义自动生成推理服务器。
在 notebook 中调用 recommend 方法试一试，看看它对某个用户给出什么推荐！在我们的实验中，模型为用户 50 预测了如下电影（图 9.12）。
图 9.12 训练后模型给出的推荐结果
对照原始数据集中的观看历史，我们会发现该用户确实偏好儿童电影和剧情片（图 9.13）。
图 9.13 用户 50 的观影示例
你可以看到，一个相当简单的模型也能给出尚可的结果，并且当它作为更大 MLOps 工作流的一部分时，还能持续自我改进。假设像我们目前这样实现了自动化数据采集与训练运行，模型会持续看到更多数据，并随着时间推移更好地泛化。当然，模型性能最终会在某个点被“封顶”，要继续提升就需要更换架构；但我们目前的基础设施搭建已经能大幅简化整个过程。事实上，Git 仓库里的代码还有一个替代模型，它在组合用户与电影 embedding 的方式上略有不同。你可以试着跑一跑，并与之前的模型做对比。
花点时间感受一下我们走了多远：你现在已经具备创建流水线的能力——处理数据、对数据做 ETL（extract, transform, load），并把清洗后的数据存入某个存储中（我们用的是 MinIO，但其他目的地也遵循同样方法）。然后我们创建了训练流水线，对模型做参数化，并用 TensorBoard 与 MLflow 为训练过程实现监控。MLflow 作为模型与实验的中央仓库，使我们能够构建复杂流水线：实现验证、推理，并通过 API 推动模型在生命周期不同阶段之间流转。我们还讨论了 VolumeOps，它通过在所有组件之间共享卷来加速训练，尤其适用于更大的数据项（例如图像）。你现在具备实现自动化流水线的技能：一端接收原始数据，另一端输出训练好的模型。这是一个相当大的里程碑——恭喜你走到这里！
我们已经探索了如何构建超越基础模型开发的稳健训练基础设施。通过使用 VolumeOps 实现高效数据处理、集成 TensorBoard 做训练可视化、以及使用 MLflow 做全面的实验追踪，我们建立了一套支持生产 ML 工作流的模式。通过目标检测与推荐系统两个例子，我们看到了这些工具如何在保持可见性与可控性的同时，简化开发流程。
系统化的实验追踪、清晰的指标体系与稳健的模型管理组合在一起，使团队能够在保持高标准生产部署的同时快速迭代。这套基础设施也构成了生产环境中模型服务的基础——我们将在下一章继续探索。下章见！
小结（Summary）
- VolumeOp 通常是训练流水线中处理大规模数据集的更优方式，它解决了每次训练运行都必须重新下载完整数据集的缺点。
- 除了模型本身，扩展的训练产物（artifacts） 同样需要追踪，包括指标、超参数集合以及验证样例。良好的血缘（lineage）管理能帮助我们在生产中更快定位与调试异常。
- MLflow 与 TensorBoard 都能追踪实验与训练运行。虽然看起来做的事情相似，但 MLflow 更适合用于跨多次实验的追踪与产物存储，而 TensorBoard 更擅长对单次运行进行更深入的“内省式”分析。
- 像 MLflow 这样集中的模型与实验追踪系统，除了让开发者工作流更顺畅之外，也能带来组织层面的收益：一个统一视窗能让干系人、项目管理与开发团队在模型能力与性能上保持同步。
- 评估指标 是模型开发的关键，因为它们用于裁决模型表现。指标也可能在模型生命周期中不断演进，因此评估框架必须支持迭代式改进。
- 对模型进行人工评估很重要。即便优秀的指标能预测性能，在早期阶段手动调试模型行为通常更快；随后再在此基础上完善指标，使其更能代表真实表现。这也意味着模型注册表必须易于访问、模型下载应足够简单。像 MLflow 这样提供下载模型 API 的注册表，在这方面非常有帮助。
