AI 容器镜像是火山引擎提供的一系列 Docker 镜像,其中预装了常见的 AI 框架、工具以及相关依赖,为用户提供快速启动且高度一致的 AI 开发环境。本文为您介绍 AI 容器镜像的使用方式。
在 AI 模型训练和推理的前置流程中,用户通常基于特定的AI框架和工具的组合开展工作(例如:PyTorch 框架、deepspeed 库),在基础环境中进行算法调优、模型注入或者脚本开发后,将重新打包好的镜像提交到训练或者推理平台的任务队列中。
AI 容器镜像用于加速用户 AI 开发、训练和推理环境的搭建,并与 AI 套件中其他工具和产品能力共同为用户提供覆盖 AI 应用完整生命周期的服务。
已经基于 VKE 搭建 AI 训练或推理任务的基础环境,例如:集群、VPC、GPU 节点等。
VKE AI 容器镜像限定只能在 VPC 内拉取镜像,暂不支持公网拉取。请在火山引擎对应地域的 VKE 集群或者 ECS 实例中拉取。
火山引擎预置提供一系列 AI 容器镜像,为用户提供快速启动且高度一致的AI开发环境。以下是 AI 容器镜像支持的版本、支持状态以及镜像地址。
说明
机器学习框架版本 | 其它软件&依赖版本 | 操作系统版本 | 是否仍在支持中 | 镜像地址 |
---|---|---|---|---|
- |
| Ubuntu 22.04 | 是 | ai-containers-{地域ID}.cr.volces.com/deeplearning/cuda11.8.0-py3.10-ubuntu22.04:latest |
Pytorch 2.0.1 |
| Ubuntu 22.04 | 是 | ai-containers-{地域ID}.cr.volces.com/deeplearning/torch2.0.1-cuda11.8.0-py3.10-ubuntu22.04:latest |
Tensorflow 2.13.0 |
| Ubuntu 22.04 | 是 | ai-containers-{地域ID}.cr.volces.com/deeplearning/tensorflow2.13.0-cuda11.8.0-py3.10-ubuntu22.04:latest |
Pytorch 2.0.1 |
| Ubuntu 22.04 | 是 | ai-containers-{地域ID}.cr.volces.com/deeplearning/megatron23.05-torch2.0.1-cuda11.8.0-py3.10-ubuntu22.04:latest |
Pytorch 2.0.1 |
| Ubuntu 22.04 | 是 | ai-containers-{地域ID}.cr.volces.com/deeplearning/deepspeed0.10.0-torch2.0.1-cuda11.8.0-py3.10-ubuntu22.04:latest |
按需获取 AI 预置容器镜像,并拉取到 VKE 中启动容器镜像,使用此 AI 容器镜像创建基础训练环境。
apiVersion: v1 kind: Pod metadata: name: ai-env # 自定义 Pod 名称。 namespace: default # 按需配置 Pod 所属的命名空间。 spec: containers: - name: pytorch # 自定义容器名称。 image: ai-containers-cn-beijing.cr.volces.com/deeplearning/torch2.0.1-cuda11.8.0-py3.10-ubuntu22.04:latest # AI 预置镜像及其版本。从步骤一中按需获取预置镜像地址,镜像版本使用 latest 即可。 command: [ "/bin/bash", "-c", "--" ] args: [ "while true; do sleep 3600; done;" ] resources: limits: nvidia.com/gpu: "1" # 按需配置 GPU 显卡数量。
Pod 进入 running 状态后,即可通过 VKE 登录 Pod 控制台,进一步进行环境准备和开发工作,具体内容可根据实际使用场景进行开发。例如:通过开发工具进行算法、脚本开发/置入模型文件。
# 所使用预置镜像的地址和版本,与步骤二中拉取的 AI 预置镜像保持一致。 FROM ai-containers-cn-beijing.cr.volces.com/deeplearning/torch2.0.1-cuda11.8.0-py3.10-ubuntu22.04:latest # 为 Dockerfile 中后续命令设置工作目录。 WORKDIR /home # 从本地文件系统中导入文件到镜像,例如:模型文件。 COPY chatglm2-6b.tar /chatglm2-6b.tar # 为镜像预装其他所需的 AI 工具软件,例如:jupyterlab。 RUN pip install "jupyterlab" # 为镜像环境添加其他所需的环境变量。 ARG VARNAME = "value" ENV PATH=$VARNAME
至此,以 AI 容器镜像为基础打包的派生镜像已经推送到火山引擎镜像仓库,您可通过 VKE 部署模型训练或推理任务。