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

AWS Deep Learning Containers（DLC）是亚马逊云科技提供的预构建深度学习容器镜像，它集成了主流深度学习框架、依赖库和工具，使开发者能够快速部署和运行深度学习工作负载。这些容器镜像经过优化，可直接在Amazon EC2、Amazon ECS、Amazon EKS和SageMaker等服务上使用。

近日，AWS DLC项目发布了基于PyTorch 2.6.0框架的训练镜像更新，支持Python 3.12环境，为开发者提供了最新的PyTorch训练能力。本次更新包含CPU和GPU两个版本，均基于Ubuntu 22.04操作系统构建。

镜像版本详情

CPU版本镜像

CPU版本镜像（pytorch-training:2.6.0-cpu-py312-ubuntu22.04-sagemaker-v1.5）主要面向不需要GPU加速的训练场景。该镜像包含了PyTorch 2.6.0的CPU版本及其相关生态组件：

• 核心框架：PyTorch 2.6.0+cpu、TorchVision 0.21.0+cpu、TorchAudio 2.6.0+cpu
• 数据处理：Pandas 2.2.3、NumPy 1.26.4、OpenCV 4.11.0.86
• 机器学习工具：scikit-learn 1.6.1、scipy 1.15.2
• AWS集成：boto3 1.37.8、sagemaker 2.241.0
• 开发工具：Cython 3.0.12、pybind11 2.13.6

GPU版本镜像

GPU版本镜像（pytorch-training:2.6.0-gpu-py312-cu126-ubuntu22.04-sagemaker-v1.5）针对CUDA 12.6环境进行了优化，适合需要GPU加速的训练任务：

• 核心框架：PyTorch 2.6.0+cu126、TorchVision 0.21.0+cu126、TorchAudio 2.6.0+cu126
• CUDA支持：包含CUDA 12.6命令行工具、cuBLAS 12.6、cuDNN等核心库
• 其他组件与CPU版本保持一致，确保开发体验的一致性

技术亮点

1. PyTorch 2.6.0新特性：该版本带来了性能优化、新算子支持以及稳定性改进，特别是在Transformer模型训练方面有显著提升。

2. Python 3.12支持：作为首个支持Python 3.12的PyTorch DLC版本，开发者可以享受到Python最新版本的语言特性和性能改进。

3. Ubuntu 22.04基础：基于长期支持的Ubuntu 22.04系统构建，提供稳定的运行环境和安全更新。

4. 完整的ML工具链：预装了从数据处理到模型训练的全套工具，包括Pandas、NumPy、scikit-learn等，开箱即用。

5. AWS服务深度集成：内置了SageMaker SDK和AWS CLI工具，方便与AWS云服务无缝集成。

使用场景

这些预构建的DLC镜像特别适合以下场景：

• 快速启动PyTorch训练任务，无需自行配置复杂的环境
• 在SageMaker上运行分布式训练作业
• 构建可重复的机器学习流水线
• 需要频繁切换不同PyTorch版本的开发环境
• 企业级机器学习平台的标准化部署

总结

AWS Deep Learning Containers的这次更新为PyTorch开发者提供了最新的2.6.0框架支持，同时保持了与AWS生态系统的紧密集成。无论是CPU还是GPU训练场景，这些预构建的容器镜像都能显著降低环境配置的复杂度，让开发者更专注于模型本身的开发和优化。对于需要在云上运行PyTorch工作负载的团队来说，这些DLC镜像是一个高效、可靠的选择。