如果您一直在阅读我们的Kubeflow系列文章，那么您现在应该知道了如何在本地安装Kubeflow和如何使用Terraform在AWS上设置Kubeflow。如果你错过了这些，去看看吧。这篇文章将等着你。的承诺。

在这篇文章中，我们将看看如何在Kubeflow集群上训练一个模型!这可能就是你一直在等待的，对吧?

培训架构会是什么样子?

架构概述如下:

在这个项目中，我们将训练一个迷你版的谷歌Imagen模型，巧妙地命名Minimagen。此外，我们将使用LAION-Aesthetics V2数据集，一些用于训练的相同数据稳定的扩散。

由于我们在Kubeflow集群上进行训练，我们必须考虑如何将数据输入和取出集群。这就是我们多年来为DagsHub添加的一些功能和集成的由来。

我们从LAION-Aesthetics V2数据集中下载了美学得分为6.5或更高的图像，并将它们放入一个公共存储库，任何人都可以访问。这个仓库有54万张图片，超过87GB !为了避免在开始之前下载整个数据集的前期成本，我们可以使用DagsHub的直接数据访问(Direct Data Access, DDA)将数据从repo流式传输到Kubeflow集群。

此外，我们可以将训练运行的参数、指标和工件记录到MLflow中。

Kubeflow Pod和DagsHub之间的所有通信都将使用dagshub和mlflowPython库!很简单，柠檬汁。

项目详细信息

我们培训项目的核心就在这里回购。它包含用于创建Docker镜像的所有内容，该镜像将在Kubeflow集群中的Pod中启动。虽然您不需要做任何特别的事情，但是对于这个repo，了解所有东西是如何组合在一起的是很重要的。这样，在将来的项目中实现它时，您将有一个良好的开端。

train.py

让我们从培训脚本。如果你看一下文件的顶部，你会看到一些环境变量是必需的:

DAGSHUB_TOKEN = os.environ。get('DAGSHUB_TOKEN'， None) DAGSHUB_USER_NAME = os.environ。get('DAGSHUB_USER_NAME'， None) DAGSHUB_REPO_NAME = os.environ。get('DAGSHUB_REPO_NAME'， None)如果DAGSHUB_TOKEN为None:引发EnvironmentError("环境变量'DAGSHUB_TOKEN'必须设置有效的令牌")如果DAGSHUB_USER_NAME为None:引发EnvironmentError("环境变量'DAGSHUB_USER_NAME'必须设置")如果DAGSHUB_REPO_NAME为None:引发EnvironmentError("环境变量'DAGSHUB_REPO_NAME'必须设置")

这样我们就可以在您的repo和运行培训的Pod之间建立通信。稍后我们将展示如何指定这些环境变量。

接下来，我们设置DagsHub数据集仓库，我们将从:

导入dagshub # 1 dagshub.auth.add_app_token(DAGSHUB_TOKEN) # 2流式导入install_hooks install_hooks(project_root='。'， repo_url='//www.kkolawyers.com/DagsHub-Datasets/LAION-Aesthetics-V2-6.5plus'， branch='main')

我们在这里:

1. 使用我们通过环境变量提供的令牌验证到DagsHub
2. 通过安装一些基于LAION数据集仓库的Python钩子来设置DDA。这些钩子检测对Python内置文件操作的调用。如果文件存在于DagsHub仓库中，它将在请求时动态加载它们。否则，将回退到本地文件系统。整洁!

之后，我们需要建立与repo相关的MLflow服务器的连接:

os.environ [' MLFLOW_TRACKING_URI '] = f“//www.kkolawyers.com/ {DAGSHUB_USER_NAME} / {DAGSHUB_REPO_NAME} .mlflow”操作系统。environ['MLFLOW_TRACKING_USERNAME'] = DAGSHUB_USER_NAME os。environ['MLFLOW_TRACKING_PASSWORD'] = DAGSHUB_TOKEN import mlflow

这允许我们使用诸如mlflow.log_params，mlflow.log_metrics,mlflow.log_artifact将内容记录到MLflow。事实上，这正是训练脚本底部的以下代码所做的:

# 1 mlflow.start_run(experiment_id=experiment_id): # 2 mlflow.log_params(args.__dict__) # 3 mlflow.log_params({"unet0_params": unets_params[0]， "unet1_params": unets_params[1]， "imagen_params": imagen_params}) # 4 mlflow.log_param("model_size_MB"， model_size_MB) # 5 mlflow.log_metrics({"unet0_avg_loss": avg_losses[0]， "unet1_avg_loss": avg_losses[1]，}) # 6 mlflow.log_artifact(dir_path)

训练完成后，这段代码:

1. 获取或创建基于名称的新MLflow实验“minimagen”
2. 记录用于运行训练脚本的参数
3. 记录模型不同部分的参数
4. 以mb为单位记录模型的大小，如脚本前面计算的那样
5. 记录两个unet的平均损失
6. 最后记录输出目录，其中包括训练好的模型

脚本的其余部分改编自最初的Minimagen repo。唯一的主要区别是a的使用自定义数据加载程序，由于数据集的格式，这是必需的。

Dockerfile

的Dockerfile对于这个项目来说是比较短的:

# 1从pytorch/pytorch:2.0.0-cuda11.7-cudnn8-runtime # 2从WORKDIR /opt/ WORKDIR拷贝/opt/workdir # 3运行apt-get update\\ g++ \\ git # 4运行pip install——no-cache-dir——upgrade -r /opt/workdir/requirements.txt # 5运行pip install——no-cache-dir——upgrade torch torchvision # 6 ENTRYPOINT ["python"， "train.py"， "——PARAMETERS"， "./ PARAMETERS"]

总之，它:

1. 从包含CUDA支持的PyTorch图像开始，这样我们就可以使用NVIDIA gpu进行训练
2. 将文件从repo复制到工作目录，这是定义的
3. 安装依赖项，这是我们需要安装的至少一个Python库所必需的
4. 安装项目所需的Python库
5. 重新安装PyTorch和TorchVision，因为步骤#4中的要求之一决定不必要地降级我们选择的版本
6. 定义Docker映像的入口点，该映像正在运行我们的训练脚本。在副本中定义了各种模型参数参数文件夹中。

说到Python需求，请查看让要查看我们需要直接安装什么:

Dagshub minimagen我流枕头

我们已经讲过大部分了。minimagen是最初的Minimagen仓库。枕头是一个功能强大的图像处理库，经常与PyTorch和ML repos一起使用。

在本地测试

要基于这个repo构建你自己的Docker镜像，你可以使用下面的命令，从最顶层的repo路径运行:

Docker build -t minimagen。

一旦它完成构建，你可以用下面的命令测试它:

\\ -e DAGSHUB_REPO_NAME= \\ -e DAGSHUB_REPO_NAME= \\ minimagen:latest——BATCH_SIZE 2——TIMESTEPS 25——TESTING

这里我们:

• 将运行容器的名称设置为minimagen-test
• (可选)启用gpu进行训练
• 传入我们提到的三个环境变量train.py脚本的需要
• 指定我们要启动标记为的Docker映像minimagen:最新
• 并向我们的训练脚本传递额外的参数

Docker image标签之后的所有内容都被附加到入口点章节。所以在这种情况下，Docker容器在启动时运行的命令是这样的:

python train.py——PARAMETERS ./ PARAMETERS \\——BATCH_SIZE 2 \\——TIMESTEPS 25 \\——TESTING .py

当一切正常工作时，你的输出应该如下所示:

2023-05-17 13:11:31.310 -信息培训尺寸:8 2023-05-17 13:11:31.310 -信息验证尺寸:8 0(00:00,它/ s ] -------------------- 时代1  -------------------- ---------- 培训……---------- 下载:100% |██████████| 773 k和773 k (0 < 00:00, 1.96 mb / s)下载:100% |██████████| 2.27 k / 2.27 k (0 < 00:00, 2.03 mb / s)下载:100% |██████████| 1.18 k / 1.18 k (0 < 00:00, 1.47 mb / s)下载:100% |██████████| 231 m / 231 m [00:02 < 00:00, 101 mb / s]一些模型的权重检查点t5-small时没有使用…-如果你从…初始化T5EncoderModel，这是预期的。如果你从…初始化T5EncoderModel，这是不可能的。T5EncoderModel的一些权重没有从模型初始化…您可能应该在一个下游任务上训练这个模型，以便能够……---------- 验证……---------- t / s] 100% |██████████| 4/4 (00:04 < 00:00,1.15 s /它]4 [00:17,4.46 s /] 00:04 < 00:00, 1.18 s /它)0 (00:00,? / s) Unet 0 avg验证失:张量(1.0752,设备= cuda: 0) Unet 1 avg验证失:张量(1.0622,设备= ' cuda: 0 ') -------------------- 时代2  -------------------- ---------- 培训……---------- ---------- 验证……---------- 100% |██████████| 4/4 (0 < 00:00,18.77 / s) 4 (00:00, 5.72 / s) 00:00 < 00:00, 21.44 / s) Unet 0 avg验证失:张量(1.0575,设备= cuda: 0) Unet 1 avg验证失:张量(1.0436,设备= cuda: 0) 100% |██████████| 4/4 (0 < 00:00,17.79 / s)

如何在Kubeflow集群上训练模型?

现在是有趣的部分!我们假设在某个地方运行着一个Kubeflow集群。如果您需要在AWS上设置一个，请查看我们之前的帖子。

由于Kubeflow构建在Kubernetes之上，您能猜到部署我们的训练映像需要什么吗?这是正确的!YAML !

在控制Kubeflow集群的机器上，创建一个新的YAML文件并调用它minimagen.yaml。将以下内容添加到文件中:

apiVersion: "kubeflow.org/v1" kind: PyTorchJob metadata: name: minimagen-train namespace: kubeflow spec: pytorchReplicaSpecs: Master: replicas: 1 restartPolicy: OnFailure template: spec: containers: - name: pytorch image: public.ecr. net . netaws/j0t5n4n3/dagshub/minimagen-train imagePullPolicy: Always args:["——BATCH_SIZE"， "2"， "——TIMESTEPS"， "25"， "——TESTING"] env: - name: DAGSHUB_TOKEN值:"" - name: DAGSHUB_USER_NAME值:"" - name: DAGSHUB_REPO_NAME值:""资源:限制:nvidia.com/gpu:

这里需要注意的重要部分是:

• 种类设为PyTorchJob。PyTorchJob是Kubernetes吗自定义资源在Kubernetes上运行PyTorch培训作业。
• 我们使用主标签下的pytorchReplicaSpecs，因为我们没有进行分布式训练。如果你在做分布式训练，你也会有工人标签下的pytorchReplicaSpecs。
• arg游戏是我们传递附加到Docker镜像的附加参数的地方入口点
• env在哪里指定我们需要的环境变量
• 资源我们在哪里告诉Kubeflow这个Docker容器需要一个GPU来运行

一旦你准备好了YAML文件，你就可以使用以下命令开始在Kubeflow集群上进行训练:

Kubectl应用-f minimage .yaml

这将使用YAML文件中指定的参数调度一个新的pod来运行您的Docker容器。

要确保一切正常，运行:

Kubectl get pods -n kubeflow | grep minimagen

你应该看到这样的内容:

minimagen-train-master-0 2/2运行0 3小时36分钟

如何删除训练仓?

如果出于某种原因需要停止pod，可以运行以下命令:

Kubectl delete -f minimage .yaml

从这里我们能去哪里?

这篇文章，连同之前的文章，已经向你展示了旋转Kubeflow集群和在单个GPU节点上开始训练的基础知识。然而,真正的当您需要跨多个节点执行分布式训练时，Kubeflow的强大功能就显现出来了。看看这个YAML例子如何开始分布式训练。

你可以做的其他改进包括:

• 增加Minimagen的大小-作为一个起点，看看。/参数文件夹，用于更改内部Minimagen模型的内部参数
• 将模型注册到MLflow-创建一个mlflow。PyFunc包装经过训练的Minimagen模型，并将其注册到MLflow模型注册表。这将允许您轻松地部署模型。有关如何做到这一点的参考，请查看我们的MLflow模型注册速成班。
• 推理-一旦你已经注册了训练模型到MLflow注册表，那么下一步就是尝试部署它!您可以通过Kubeflow或其他方式完成此操作。可以找到关于设置推理服务的更多信息在这里和在这里。
• 记录实时训练指标到Mlflow-目前，培训指标只在培训结束时记录。能够看到训练指标在训练过程中是如何演变的，这可能会更有趣(也更有用!)。

这些改进中的每一个都是repo上的开放问题。欢迎pr !如果你需要关于如何开始的帮助或建议，请随时提问

---

恭喜你!通过本系列文章，您已经了解了很多关于设置Kubeflow和使用它来训练模型的知识。你值得我们击掌庆祝!

Kubeflow还有很多东西等着你去发现。看看你可以如何使用笔记本电脑和设置管道首先!

如果你有任何问题，请随时联系我。你可以加入我们的不和在美国，我们建立了一个充满活力、乐于助人、友好的社区。

标签

• Kubeflow
• MLOps
• 教程
• 数据科学工作流程