如何在MLflow中完善数据集列？（猫狗分类Python应用场景）

用MLflow的Dataset API就能搞定，具体步骤如下： #

1. 先导入所需依赖 #

import mlflow
from mlflow.data import Dataset
from mlflow.data.local_source import LocalSource

2. 定义你的猫狗数据集信息 #

假设数据集存在本地./cat_dog_dataset路径下，直接把数据集的关键信息封装成Dataset对象：

dataset_path = "./cat_dog_dataset"
# 指定数据源
source = LocalSource(dataset_path)
# 填充数据集元数据：名称、描述、分类标签等
dataset = Dataset(
    source=source,
    name="猫狗二分类数据集",
    description="包含训练/验证/测试三部分，用于猫狗图像分类任务",
    tags={
        "任务类型": "二分类",
        "类别数量": 2,
        "训练集样本数": 1000,
        "验证集样本数": 200,
        "测试集样本数": 200,
        "图像尺寸": "224x224",
        "图像格式": "JPEG"
    }
)

3. 在训练代码中记录数据集 #

启动MLflow运行时，用mlflow.log_input()把数据集信息上传：

with mlflow.start_run():
    # 记录数据集到当前运行
    mlflow.log_input(dataset, context="training")
    # 这里放你的猫狗分类训练代码...

4. 查看完善后的效果 #

运行代码后打开MLflow UI，进入对应实验的数据集标签页，就能看到你补充的所有信息，对应你截图里的数据集列内容。

内容的提问来源于stack exchange，提问作者skyumin09