AWS Deep Learning Containers发布PyTorch 2.6.0训练镜像

AWS Deep Learning Containers（DLC）项目为机器学习开发者提供了预配置的深度学习环境容器镜像，这些镜像经过AWS官方优化，可直接用于模型训练和推理任务。该项目支持主流深度学习框架如PyTorch、TensorFlow等，并针对AWS EC2实例进行了性能优化。

近日，AWS DLC项目发布了基于PyTorch 2.6.0框架的新版本训练镜像，支持Python 3.12环境。这些镜像分为CPU和GPU两个版本，均基于Ubuntu 22.04操作系统构建。

CPU版本镜像特性

CPU版本镜像（pytorch-training:2.6.0-cpu-py312-ubuntu22.04-ec2）主要包含以下技术栈：

• 基础环境：Ubuntu 22.04操作系统
• Python版本：3.12
• PyTorch版本：2.6.0（CPU优化版）
• 关键科学计算库：
  • NumPy 2.2.3
  • SciPy 1.15.2
  • OpenCV 4.11.0
  • Cython 3.0.12
  • MKL 2025.0.1（Intel数学核心库）
• 工具链：
  • AWS CLI 1.38.8
  • Boto3 1.37.8
  • MPI4py 4.0.3（支持分布式训练）

该镜像特别适合不需要GPU加速的轻量级模型训练任务，或者作为开发测试环境使用。

GPU版本镜像特性

GPU版本镜像（pytorch-training:2.6.0-gpu-py312-cu126-ubuntu22.04-ec2）在CPU版本基础上增加了对CUDA 12.6的支持，包含：

• CUDA工具链：12.6版本
• cuDNN：9.x版本（针对CUDA 12优化）
• NCCL：高性能GPU通信库
• PyTorch GPU版本：2.6.0+cu126
• TorchVision：0.21.0+cu126
• TorchAudio：2.6.0+cu126

GPU版本镜像充分利用了NVIDIA GPU的并行计算能力，适合大规模深度学习模型训练任务。镜像中包含了必要的GPU驱动和库文件，开箱即用。

技术亮点

1. Python 3.12支持：新版本镜像率先支持Python 3.12，开发者可以使用最新的Python特性。

2. PyTorch 2.6.0优化：PyTorch 2.6带来了多项性能改进和新特性，包括改进的编译器后端、更高效的算子实现等。

3. 科学计算栈更新：NumPy 2.x系列带来了显著的性能提升，SciPy 1.15.2提供了更稳定的科学计算功能。

4. 构建系统优化：使用Ninja 1.11.1作为构建系统，显著加快了大型项目的编译速度。

5. 开发工具集成：镜像中预装了Emacs等开发工具，方便开发者直接在容器中进行代码编辑和调试。

使用建议

对于选择使用这些镜像的开发者，AWS提供以下建议：

1. 根据计算需求选择合适的版本：轻量级任务使用CPU版本即可，大规模训练任务建议使用GPU版本。

2. 利用预装的AWS CLI和Boto3工具可以方便地与AWS服务交互，如从S3存取训练数据。

3. 对于分布式训练场景，MPI4py和NCCL库已经集成，可以直接使用。

4. 镜像中的科学计算栈已经过兼容性测试，建议优先使用预装版本，避免自行安装可能导致的依赖冲突。

这些经过AWS优化的DLC镜像大大简化了深度学习环境的搭建过程，开发者可以专注于模型开发而非环境配置，显著提高工作效率。