chaneyk/pytorch

PyTorch GPU支持镜像

镜像概述

本镜像基于daniilidis-group-base镜像构建，提供支持GPU加速的PyTorch版本，专为深度学习任务设计。镜像包含Python 2和Python 3两种版本，可满足不同项目的开发环境需求，确保在GPU硬件上高效运行PyTorch应用。

核心功能与特性

• GPU加速支持：针对NVIDIA GPU进行优化，充分利用显卡计算能力加速深度学习模型训练与推理
• 多Python版本兼容：同时提供Python 2和Python 3版本，适配不同项目的依赖需求
• 硬件兼容性：支持多种主流NVIDIA GPU型号
• 稳定依赖配置：绑定特定版本的CUDA和cuDNN，确保开发环境一致性和稳定性

支持的硬件

• Tesla V100
• GeForce 1080Ti
• Titan X
• Titan Z

CUDA/cuDNN版本对应关系

使用场景

• 深度学习模型的训练与推理任务
• 基于PyTorch的学术研究与实验
• 需要GPU加速的机器学习应用开发
• 多Python版本环境下的PyTorch项目测试

使用方法

拉取镜像

根据需求选择Python版本和PyTorch版本拉取对应镜像（具体标签需参考官方仓库）：

bash
# 示例：拉取Python 3版本的PyTorch 0.4.0镜像
docker pull daniilidis-group/pytorch:0.4.0-py3

# 示例：拉取Python 2版本的PyTorch 0.3.1镜像
docker pull daniilidis-group/pytorch:0.3.1-py2

运行容器

运行容器时需确保主机已安装NVIDIA Docker运行时，并添加GPU支持参数：

bash
# 使用所有GPU运行Python 3版本容器
docker run --gpus all -it daniilidis-group/pytorch:0.4.0-py3 /bin/bash

# 指定特定GPU运行容器（例如使用第0块GPU）
docker run --gpus "device=0" -it daniilidis-group/pytorch:0.3.1-py2 /bin/bash

验证环境

容器启动后，可通过以下命令验证PyTorch和GPU是否正常工作：

python
import torch
print("PyTorch版本:", torch.__version__)
print("CUDA是否可用:", torch.cuda.is_available())
print("GPU设备数量:", torch.cuda.device_count())
print("当前GPU名称:", torch.cuda.get_device_name(0))