小白也能轻松上手PyTorch 2.5镜像详细使用教程涵盖从安装到运行你是不是刚接触深度学习被各种环境配置搞得头大或者你是个有经验的开发者但每次在新机器上搭PyTorch环境都要花半天时间折腾CUDA版本、驱动兼容这些破事如果你点头了那今天这个教程就是为你准备的。我要带你体验一种完全不同的方式——用PyTorch 2.5镜像10分钟搞定从零到能跑代码的完整环境。简单说这个镜像就像个“即食套餐”。别人还在买菜、洗菜、切菜、开火你已经吃上了。它里面预装了PyTorch 2.5、torchvision、torchaudio还有匹配好的CUDA全都调好了味道你打开就能用。这篇文章我会手把手带你走一遍完整流程怎么获取镜像、怎么启动、怎么用两种方式操作最后再跑几个实际例子验证环境。跟着做一遍你就能彻底摆脱环境配置的噩梦。1. 准备工作你只需要这三样东西在开始之前我们先看看需要准备什么。其实特别简单就三样1. 一台有NVIDIA显卡的电脑笔记本或台式机都行显卡最好是GTX 10系列或更新比如RTX 20/30/40系列如果没有独立显卡也能用只是训练速度会慢很多2. 安装好DockerDocker是个容器工具你可以把它理解成一个“轻量级虚拟机”去Docker官网下载对应你系统的版本Windows/Mac/Linux安装过程就是一路下一步没什么难度3. 确保NVIDIA驱动正常右键点击桌面如果有“NVIDIA控制面板”就说明驱动装好了如果没有去NVIDIA官网下载对应你显卡的驱动安装就这些没了。不需要你自己装CUDA不需要配环境变量不需要担心版本冲突。所有的复杂工作镜像都帮你做好了。2. 获取和启动镜像两步搞定环境准备好了吗我们现在开始实际操作。整个过程就两步拉取镜像、启动容器。2.1 第一步拉取PyTorch 2.5镜像打开你的终端Windows用PowerShell或CMDMac/Linux用Terminal输入下面这个命令docker pull pytorch/pytorch:2.5.0-cuda12.4-cudnn9-runtime这个命令是什么意思呢我解释一下docker pull告诉Docker去下载镜像pytorch/pytorch这是镜像在Docker Hub上的名字2.5.0-cuda12.4-cudnn9-runtime这是标签指定了版本2.5.0PyTorch主版本cuda12.4对应的CUDA版本cudnn9对应的cuDNN版本深度学习加速库runtime运行环境版本比较轻量等它下载完你会看到类似这样的提示Status: Downloaded newer image for pytorch/pytorch:2.5.0-cuda12.4-cudnn9-runtime这就表示镜像下载成功了。整个过程取决于你的网速一般几分钟到十几分钟。2.2 第二步启动容器两种方式任选镜像下载好了但它还是个“静态”的文件。我们需要把它运行起来变成“动态”的容器。这里我给你两种启动方式你可以根据喜好选择。方式一用Jupyter Notebook启动推荐新手Jupyter Notebook是个网页版的编程环境特别适合学习和实验。你可以一边写代码一边看结果像记笔记一样。用这个命令启动docker run -it --gpus all -p 8888:8888 -v /本地/你的/目录:/workspace pytorch/pytorch:2.5.0-cuda12.4-cudnn9-runtime jupyter notebook --ip0.0.0.0 --allow-root --NotebookApp.token别怕这个命令看起来长我拆开给你解释docker run运行容器-it交互模式让你能输入命令--gpus all使用所有GPU如果没有GPU去掉这个参数-p 8888:8888把容器的8888端口映射到本机的8888端口-v /本地/你的/目录:/workspace把本地目录挂载到容器的/workspace目录这样文件可以互通最后是启动Jupyter Notebook的命令启动成功后终端会显示一个链接类似http://127.0.0.1:8888/?token...用浏览器打开这个链接你就进入了Jupyter界面。方式二用命令行启动推荐开发者如果你习惯在终端里工作或者要跑长时间的训练任务可以用这种方式docker run -it --gpus all -p 2222:22 -v /本地/你的/目录:/workspace pytorch/pytorch:2.5.0-cuda12.4-cudnn9-runtime /bin/bash这个命令会直接进入容器的命令行。进去之后你就像在一台新的Linux电脑里一样可以执行各种命令。3. 验证环境确保一切正常环境启动起来了我们得确认一下是不是真的能用。特别是GPU能不能用这是最重要的。3.1 检查PyTorch和CUDA不管你是用Jupyter还是命令行先创建一个Python环境检查一下。如果你在Jupyter里新建一个Notebook如果在命令行输入python进入Python交互模式。然后输入这几行代码import torch # 打印PyTorch版本 print(PyTorch版本:, torch.__version__) # 检查CUDA是否可用 print(CUDA是否可用:, torch.cuda.is_available()) # 如果有CUDA打印显卡信息 if torch.cuda.is_available(): print(显卡名称:, torch.cuda.get_device_name(0)) print(CUDA版本:, torch.version.cuda)你应该看到类似这样的输出PyTorch版本: 2.5.0 CUDA是否可用: True 显卡名称: NVIDIA GeForce RTX 4090 CUDA版本: 12.4如果CUDA是否可用显示True恭喜你GPU加速已经就绪了。如果显示False可能是你的Docker没配置好GPU支持或者显卡驱动有问题。3.2 检查torchvision和torchaudio这两个是PyTorch的扩展库一个处理图像一个处理音频。检查一下它们是不是也装好了import torchvision import torchaudio print(torchvision版本:, torchvision.__version__) print(torchaudio版本:, torchaudio.__version__)正常的话你会看到它们的版本号说明都预装好了。4. 实际动手跑几个例子感受一下环境验证通过了我们实际用一下看看这个镜像到底能做什么。我给你准备三个小例子从简单到稍微复杂一点。4.1 例子一用GPU加速计算30秒第一个例子我们看看GPU能带来多快的速度提升。创建一个新的代码单元格输入import torch import time # 创建一个大的随机矩阵 size 5000 cpu_tensor torch.randn(size, size) # 在CPU上做矩阵乘法 start_time time.time() cpu_result cpu_tensor cpu_tensor cpu_time time.time() - start_time print(fCPU计算时间: {cpu_time:.2f}秒) # 把数据移到GPU上 if torch.cuda.is_available(): gpu_tensor cpu_tensor.cuda() # 在GPU上做同样的计算 start_time time.time() gpu_result gpu_tensor gpu_tensor torch.cuda.synchronize() # 等待GPU计算完成 gpu_time time.time() - start_time print(fGPU计算时间: {gpu_time:.2f}秒) print(f加速比: {cpu_time/gpu_time:.1f}倍)运行这个代码你会看到GPU比CPU快多少倍。在我的测试里RTX 4090比i9 CPU快了大概50倍。这个差距在训练大模型时会非常明显。4.2 例子二加载预训练模型做图像分类2分钟torchvision里有很多现成的模型我们直接用它们来做点实际的事情。这个例子我们用ResNet-50模型识别一张图片import torch import torchvision import torchvision.transforms as transforms from PIL import Image import requests from io import BytesIO # 1. 加载预训练的ResNet-50模型 model torchvision.models.resnet50(pretrainedTrue) model.eval() # 设置为评估模式 # 2. 准备图像预处理 preprocess transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]), ]) # 3. 下载一张测试图片这里用一只猫的图片 url https://upload.wikimedia.org/wikipedia/commons/thumb/3/3a/Cat03.jpg/1200px-Cat03.jpg response requests.get(url) img Image.open(BytesIO(response.content)) # 4. 预处理图片 input_tensor preprocess(img) input_batch input_tensor.unsqueeze(0) # 增加一个批次维度 # 5. 如果有GPU把数据移到GPU上 if torch.cuda.is_available(): input_batch input_batch.to(cuda) model.to(cuda) # 6. 运行模型 with torch.no_grad(): # 不计算梯度节省内存 output model(input_batch) # 7. 读取分类结果 # 加载ImageNet的类别标签 import json url https://raw.githubusercontent.com/anishathalye/imagenet-simple-labels/master/imagenet-simple-labels.json labels json.loads(requests.get(url).text) # 获取前5个最可能的类别 probabilities torch.nn.functional.softmax(output[0], dim0) top5_prob, top5_catid torch.topk(probabilities, 5) print(识别结果前5个可能类别) for i in range(top5_prob.size(0)): print(f{labels[top5_catid[i]]}: {top5_prob[i].item()*100:.2f}%)运行这个代码模型会分析图片然后告诉你它认为图片里是什么。ResNet-50是在ImageNet数据集上训练的能识别1000种常见物体。4.3 例子三简单的神经网络训练5分钟我们从头训练一个小的神经网络体验一下完整的流程import torch import torch.nn as nn import torch.optim as optim import torchvision import torchvision.transforms as transforms from torch.utils.data import DataLoader # 1. 准备数据用MNIST手写数字数据集 transform transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,)) ]) train_dataset torchvision.datasets.MNIST( root./data, trainTrue, downloadTrue, transformtransform ) test_dataset torchvision.datasets.MNIST( root./data, trainFalse, downloadTrue, transformtransform ) # 2. 创建数据加载器 batch_size 64 train_loader DataLoader(train_dataset, batch_sizebatch_size, shuffleTrue) test_loader DataLoader(test_dataset, batch_sizebatch_size, shuffleFalse) # 3. 定义一个简单的神经网络 class SimpleNN(nn.Module): def __init__(self): super(SimpleNN, self).__init__() self.flatten nn.Flatten() self.network nn.Sequential( nn.Linear(28*28, 128), nn.ReLU(), nn.Linear(128, 64), nn.ReLU(), nn.Linear(64, 10) ) def forward(self, x): x self.flatten(x) return self.network(x) # 4. 创建模型、损失函数和优化器 device torch.device(cuda if torch.cuda.is_available() else cpu) model SimpleNN().to(device) criterion nn.CrossEntropyLoss() optimizer optim.Adam(model.parameters(), lr0.001) # 5. 训练函数 def train(epoch): model.train() running_loss 0.0 for batch_idx, (data, target) in enumerate(train_loader): data, target data.to(device), target.to(device) optimizer.zero_grad() output model(data) loss criterion(output, target) loss.backward() optimizer.step() running_loss loss.item() if batch_idx % 100 99: print(fEpoch: {epoch}, Batch: {batch_idx1}, Loss: {running_loss/100:.3f}) running_loss 0.0 # 6. 测试函数 def test(): model.eval() correct 0 total 0 with torch.no_grad(): for data, target in test_loader: data, target data.to(device), target.to(device) output model(data) _, predicted torch.max(output.data, 1) total target.size(0) correct (predicted target).sum().item() accuracy 100 * correct / total print(f测试准确率: {accuracy:.2f}%) return accuracy # 7. 开始训练 print(开始训练...) for epoch in range(1, 6): # 训练5个epoch train(epoch) test() print(训练完成)这个例子虽然简单但包含了深度学习的核心流程准备数据、定义模型、训练、测试。你可以看到GPU如何加速训练过程。5. 常见问题与解决方案刚开始用的时候可能会遇到一些小问题。我整理了最常见的几个问题1Docker命令报错说找不到GPUdocker: Error response from daemon: could not select device driver with capabilities: [[gpu]].解决你需要安装NVIDIA Container Toolkit。具体步骤先卸载旧的nvidia-docker如果有docker volume ls -q -f drivernvidia-docker | xargs -r -I{} docker rm -f {} sudo apt-get purge -y nvidia-docker添加NVIDIA仓库distribution$(. /etc/os-release;echo $ID$VERSION_ID) curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add - curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.list | sudo tee /etc/apt/sources.list.d/nvidia-docker.list安装sudo apt-get update sudo apt-get install -y nvidia-container-toolkit sudo systemctl restart docker问题2内存不足训练时崩溃解决可以调整batch size批量大小或者使用梯度累积# 如果原来batch_size64导致内存不足可以减小 batch_size 32 # 或者使用梯度累积相当于大的batch size但内存占用小 accumulation_steps 4 # 每4个batch更新一次参数 optimizer.zero_grad() for i, (data, target) in enumerate(train_loader): # ... 前向传播和损失计算 loss loss / accumulation_steps # 标准化损失 loss.backward() if (i1) % accumulation_steps 0: optimizer.step() optimizer.zero_grad()问题3想安装额外的Python包解决在容器内直接用pip安装# 如果你在容器内通过Jupyter或SSH pip install pandas matplotlib seaborn # 或者构建自己的Dockerfile基于这个镜像 # Dockerfile内容 FROM pytorch/pytorch:2.5.0-cuda12.4-cudnn9-runtime RUN pip install pandas matplotlib seaborn问题4如何保存和加载模型解决PyTorch提供了简单的方法# 保存整个模型 torch.save(model, model.pth) # 只保存模型参数推荐更灵活 torch.save(model.state_dict(), model_weights.pth) # 加载模型 # 方式一加载整个模型 model torch.load(model.pth) # 方式二先创建模型结构再加载参数 model SimpleNN() # 需要先定义同样的模型结构 model.load_state_dict(torch.load(model_weights.pth)) model.eval() # 记得设置为评估模式6. 进阶技巧让开发更高效掌握了基础用法后我再分享几个让开发更高效的小技巧6.1 使用Docker Compose管理多个服务如果你需要同时运行多个容器比如一个跑模型训练一个跑API服务可以用Docker Compose。创建一个docker-compose.yml文件version: 3.8 services: pytorch-jupyter: image: pytorch/pytorch:2.5.0-cuda12.4-cudnn9-runtime runtime: nvidia ports: - 8888:8888 volumes: - ./workspace:/workspace command: jupyter notebook --ip0.0.0.0 --allow-root --NotebookApp.token pytorch-api: build: ./api # 假设你有个API服务的Dockerfile ports: - 5000:5000 depends_on: - pytorch-jupyter然后运行docker-compose up6.2 配置Jupyter扩展让Jupyter更好用的一些扩展# 在容器内安装 pip install jupyter_contrib_nbextensions jupyter contrib nbextension install --user # 启用有用的扩展 jupyter nbextension enable code_prettify/code_prettify jupyter nbextension enable execute_time/ExecuteTime jupyter nbextension enable toggle_all_line_numbers/main6.3 使用TensorBoard可视化训练过程PyTorch支持TensorBoard可以可视化训练过程from torch.utils.tensorboard import SummaryWriter # 创建writer writer SummaryWriter(runs/experiment_1) # 在训练循环中记录 for epoch in range(num_epochs): # ... 训练代码 writer.add_scalar(Loss/train, loss.item(), epoch) writer.add_scalar(Accuracy/train, accuracy, epoch) # 在终端启动TensorBoard # tensorboard --logdirruns6.4 利用torch.compile加速PyTorch 2.x新特性PyTorch 2.5的torch.compile可以显著加速模型import torch # 普通的模型 model SimpleNN().cuda() # 编译后的模型第一次运行会慢之后会快 compiled_model torch.compile(model) # 使用编译后的模型 output compiled_model(input_data)7. 总结走完这个完整的教程你现在应该已经成功安装并运行了PyTorch 2.5镜像- 不用再折腾环境配置了验证了GPU加速可用- 看到GPU比CPU快了多少倍实际运行了三个例子- 从简单计算到完整训练流程都体验了学会了解决常见问题- 遇到问题知道怎么查怎么解决了了解了一些进阶技巧- 可以让你的开发更高效这个镜像最大的价值就是帮你跳过了深度学习中最枯燥、最易出错的环境配置环节。你不用再担心我的CUDA版本和PyTorch版本匹配吗torchvision怎么装不上为什么检测不到GPU所有这些问题镜像都帮你解决了。你拿到的是一个完整、一致、可复现的环境。无论是自己学习还是团队协作或者是部署到生产环境都用同一个镜像保证结果一致。我建议你保存好今天用到的命令和代码以后每次开始新项目都可以快速搭建环境。深度学习应该把时间花在思考模型、调整参数、分析结果上而不是浪费在配环境上。获取更多AI镜像想探索更多AI镜像和应用场景访问 CSDN星图镜像广场提供丰富的预置镜像覆盖大模型推理、图像生成、视频生成、模型微调等多个领域支持一键部署。
小白也能轻松上手：PyTorch 2.5镜像详细使用教程，涵盖从安装到运行
发布时间：2026/5/27 21:48:56
小白也能轻松上手PyTorch 2.5镜像详细使用教程涵盖从安装到运行你是不是刚接触深度学习被各种环境配置搞得头大或者你是个有经验的开发者但每次在新机器上搭PyTorch环境都要花半天时间折腾CUDA版本、驱动兼容这些破事如果你点头了那今天这个教程就是为你准备的。我要带你体验一种完全不同的方式——用PyTorch 2.5镜像10分钟搞定从零到能跑代码的完整环境。简单说这个镜像就像个“即食套餐”。别人还在买菜、洗菜、切菜、开火你已经吃上了。它里面预装了PyTorch 2.5、torchvision、torchaudio还有匹配好的CUDA全都调好了味道你打开就能用。这篇文章我会手把手带你走一遍完整流程怎么获取镜像、怎么启动、怎么用两种方式操作最后再跑几个实际例子验证环境。跟着做一遍你就能彻底摆脱环境配置的噩梦。1. 准备工作你只需要这三样东西在开始之前我们先看看需要准备什么。其实特别简单就三样1. 一台有NVIDIA显卡的电脑笔记本或台式机都行显卡最好是GTX 10系列或更新比如RTX 20/30/40系列如果没有独立显卡也能用只是训练速度会慢很多2. 安装好DockerDocker是个容器工具你可以把它理解成一个“轻量级虚拟机”去Docker官网下载对应你系统的版本Windows/Mac/Linux安装过程就是一路下一步没什么难度3. 确保NVIDIA驱动正常右键点击桌面如果有“NVIDIA控制面板”就说明驱动装好了如果没有去NVIDIA官网下载对应你显卡的驱动安装就这些没了。不需要你自己装CUDA不需要配环境变量不需要担心版本冲突。所有的复杂工作镜像都帮你做好了。2. 获取和启动镜像两步搞定环境准备好了吗我们现在开始实际操作。整个过程就两步拉取镜像、启动容器。2.1 第一步拉取PyTorch 2.5镜像打开你的终端Windows用PowerShell或CMDMac/Linux用Terminal输入下面这个命令docker pull pytorch/pytorch:2.5.0-cuda12.4-cudnn9-runtime这个命令是什么意思呢我解释一下docker pull告诉Docker去下载镜像pytorch/pytorch这是镜像在Docker Hub上的名字2.5.0-cuda12.4-cudnn9-runtime这是标签指定了版本2.5.0PyTorch主版本cuda12.4对应的CUDA版本cudnn9对应的cuDNN版本深度学习加速库runtime运行环境版本比较轻量等它下载完你会看到类似这样的提示Status: Downloaded newer image for pytorch/pytorch:2.5.0-cuda12.4-cudnn9-runtime这就表示镜像下载成功了。整个过程取决于你的网速一般几分钟到十几分钟。2.2 第二步启动容器两种方式任选镜像下载好了但它还是个“静态”的文件。我们需要把它运行起来变成“动态”的容器。这里我给你两种启动方式你可以根据喜好选择。方式一用Jupyter Notebook启动推荐新手Jupyter Notebook是个网页版的编程环境特别适合学习和实验。你可以一边写代码一边看结果像记笔记一样。用这个命令启动docker run -it --gpus all -p 8888:8888 -v /本地/你的/目录:/workspace pytorch/pytorch:2.5.0-cuda12.4-cudnn9-runtime jupyter notebook --ip0.0.0.0 --allow-root --NotebookApp.token别怕这个命令看起来长我拆开给你解释docker run运行容器-it交互模式让你能输入命令--gpus all使用所有GPU如果没有GPU去掉这个参数-p 8888:8888把容器的8888端口映射到本机的8888端口-v /本地/你的/目录:/workspace把本地目录挂载到容器的/workspace目录这样文件可以互通最后是启动Jupyter Notebook的命令启动成功后终端会显示一个链接类似http://127.0.0.1:8888/?token...用浏览器打开这个链接你就进入了Jupyter界面。方式二用命令行启动推荐开发者如果你习惯在终端里工作或者要跑长时间的训练任务可以用这种方式docker run -it --gpus all -p 2222:22 -v /本地/你的/目录:/workspace pytorch/pytorch:2.5.0-cuda12.4-cudnn9-runtime /bin/bash这个命令会直接进入容器的命令行。进去之后你就像在一台新的Linux电脑里一样可以执行各种命令。3. 验证环境确保一切正常环境启动起来了我们得确认一下是不是真的能用。特别是GPU能不能用这是最重要的。3.1 检查PyTorch和CUDA不管你是用Jupyter还是命令行先创建一个Python环境检查一下。如果你在Jupyter里新建一个Notebook如果在命令行输入python进入Python交互模式。然后输入这几行代码import torch # 打印PyTorch版本 print(PyTorch版本:, torch.__version__) # 检查CUDA是否可用 print(CUDA是否可用:, torch.cuda.is_available()) # 如果有CUDA打印显卡信息 if torch.cuda.is_available(): print(显卡名称:, torch.cuda.get_device_name(0)) print(CUDA版本:, torch.version.cuda)你应该看到类似这样的输出PyTorch版本: 2.5.0 CUDA是否可用: True 显卡名称: NVIDIA GeForce RTX 4090 CUDA版本: 12.4如果CUDA是否可用显示True恭喜你GPU加速已经就绪了。如果显示False可能是你的Docker没配置好GPU支持或者显卡驱动有问题。3.2 检查torchvision和torchaudio这两个是PyTorch的扩展库一个处理图像一个处理音频。检查一下它们是不是也装好了import torchvision import torchaudio print(torchvision版本:, torchvision.__version__) print(torchaudio版本:, torchaudio.__version__)正常的话你会看到它们的版本号说明都预装好了。4. 实际动手跑几个例子感受一下环境验证通过了我们实际用一下看看这个镜像到底能做什么。我给你准备三个小例子从简单到稍微复杂一点。4.1 例子一用GPU加速计算30秒第一个例子我们看看GPU能带来多快的速度提升。创建一个新的代码单元格输入import torch import time # 创建一个大的随机矩阵 size 5000 cpu_tensor torch.randn(size, size) # 在CPU上做矩阵乘法 start_time time.time() cpu_result cpu_tensor cpu_tensor cpu_time time.time() - start_time print(fCPU计算时间: {cpu_time:.2f}秒) # 把数据移到GPU上 if torch.cuda.is_available(): gpu_tensor cpu_tensor.cuda() # 在GPU上做同样的计算 start_time time.time() gpu_result gpu_tensor gpu_tensor torch.cuda.synchronize() # 等待GPU计算完成 gpu_time time.time() - start_time print(fGPU计算时间: {gpu_time:.2f}秒) print(f加速比: {cpu_time/gpu_time:.1f}倍)运行这个代码你会看到GPU比CPU快多少倍。在我的测试里RTX 4090比i9 CPU快了大概50倍。这个差距在训练大模型时会非常明显。4.2 例子二加载预训练模型做图像分类2分钟torchvision里有很多现成的模型我们直接用它们来做点实际的事情。这个例子我们用ResNet-50模型识别一张图片import torch import torchvision import torchvision.transforms as transforms from PIL import Image import requests from io import BytesIO # 1. 加载预训练的ResNet-50模型 model torchvision.models.resnet50(pretrainedTrue) model.eval() # 设置为评估模式 # 2. 准备图像预处理 preprocess transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]), ]) # 3. 下载一张测试图片这里用一只猫的图片 url https://upload.wikimedia.org/wikipedia/commons/thumb/3/3a/Cat03.jpg/1200px-Cat03.jpg response requests.get(url) img Image.open(BytesIO(response.content)) # 4. 预处理图片 input_tensor preprocess(img) input_batch input_tensor.unsqueeze(0) # 增加一个批次维度 # 5. 如果有GPU把数据移到GPU上 if torch.cuda.is_available(): input_batch input_batch.to(cuda) model.to(cuda) # 6. 运行模型 with torch.no_grad(): # 不计算梯度节省内存 output model(input_batch) # 7. 读取分类结果 # 加载ImageNet的类别标签 import json url https://raw.githubusercontent.com/anishathalye/imagenet-simple-labels/master/imagenet-simple-labels.json labels json.loads(requests.get(url).text) # 获取前5个最可能的类别 probabilities torch.nn.functional.softmax(output[0], dim0) top5_prob, top5_catid torch.topk(probabilities, 5) print(识别结果前5个可能类别) for i in range(top5_prob.size(0)): print(f{labels[top5_catid[i]]}: {top5_prob[i].item()*100:.2f}%)运行这个代码模型会分析图片然后告诉你它认为图片里是什么。ResNet-50是在ImageNet数据集上训练的能识别1000种常见物体。4.3 例子三简单的神经网络训练5分钟我们从头训练一个小的神经网络体验一下完整的流程import torch import torch.nn as nn import torch.optim as optim import torchvision import torchvision.transforms as transforms from torch.utils.data import DataLoader # 1. 准备数据用MNIST手写数字数据集 transform transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,)) ]) train_dataset torchvision.datasets.MNIST( root./data, trainTrue, downloadTrue, transformtransform ) test_dataset torchvision.datasets.MNIST( root./data, trainFalse, downloadTrue, transformtransform ) # 2. 创建数据加载器 batch_size 64 train_loader DataLoader(train_dataset, batch_sizebatch_size, shuffleTrue) test_loader DataLoader(test_dataset, batch_sizebatch_size, shuffleFalse) # 3. 定义一个简单的神经网络 class SimpleNN(nn.Module): def __init__(self): super(SimpleNN, self).__init__() self.flatten nn.Flatten() self.network nn.Sequential( nn.Linear(28*28, 128), nn.ReLU(), nn.Linear(128, 64), nn.ReLU(), nn.Linear(64, 10) ) def forward(self, x): x self.flatten(x) return self.network(x) # 4. 创建模型、损失函数和优化器 device torch.device(cuda if torch.cuda.is_available() else cpu) model SimpleNN().to(device) criterion nn.CrossEntropyLoss() optimizer optim.Adam(model.parameters(), lr0.001) # 5. 训练函数 def train(epoch): model.train() running_loss 0.0 for batch_idx, (data, target) in enumerate(train_loader): data, target data.to(device), target.to(device) optimizer.zero_grad() output model(data) loss criterion(output, target) loss.backward() optimizer.step() running_loss loss.item() if batch_idx % 100 99: print(fEpoch: {epoch}, Batch: {batch_idx1}, Loss: {running_loss/100:.3f}) running_loss 0.0 # 6. 测试函数 def test(): model.eval() correct 0 total 0 with torch.no_grad(): for data, target in test_loader: data, target data.to(device), target.to(device) output model(data) _, predicted torch.max(output.data, 1) total target.size(0) correct (predicted target).sum().item() accuracy 100 * correct / total print(f测试准确率: {accuracy:.2f}%) return accuracy # 7. 开始训练 print(开始训练...) for epoch in range(1, 6): # 训练5个epoch train(epoch) test() print(训练完成)这个例子虽然简单但包含了深度学习的核心流程准备数据、定义模型、训练、测试。你可以看到GPU如何加速训练过程。5. 常见问题与解决方案刚开始用的时候可能会遇到一些小问题。我整理了最常见的几个问题1Docker命令报错说找不到GPUdocker: Error response from daemon: could not select device driver with capabilities: [[gpu]].解决你需要安装NVIDIA Container Toolkit。具体步骤先卸载旧的nvidia-docker如果有docker volume ls -q -f drivernvidia-docker | xargs -r -I{} docker rm -f {} sudo apt-get purge -y nvidia-docker添加NVIDIA仓库distribution$(. /etc/os-release;echo $ID$VERSION_ID) curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add - curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.list | sudo tee /etc/apt/sources.list.d/nvidia-docker.list安装sudo apt-get update sudo apt-get install -y nvidia-container-toolkit sudo systemctl restart docker问题2内存不足训练时崩溃解决可以调整batch size批量大小或者使用梯度累积# 如果原来batch_size64导致内存不足可以减小 batch_size 32 # 或者使用梯度累积相当于大的batch size但内存占用小 accumulation_steps 4 # 每4个batch更新一次参数 optimizer.zero_grad() for i, (data, target) in enumerate(train_loader): # ... 前向传播和损失计算 loss loss / accumulation_steps # 标准化损失 loss.backward() if (i1) % accumulation_steps 0: optimizer.step() optimizer.zero_grad()问题3想安装额外的Python包解决在容器内直接用pip安装# 如果你在容器内通过Jupyter或SSH pip install pandas matplotlib seaborn # 或者构建自己的Dockerfile基于这个镜像 # Dockerfile内容 FROM pytorch/pytorch:2.5.0-cuda12.4-cudnn9-runtime RUN pip install pandas matplotlib seaborn问题4如何保存和加载模型解决PyTorch提供了简单的方法# 保存整个模型 torch.save(model, model.pth) # 只保存模型参数推荐更灵活 torch.save(model.state_dict(), model_weights.pth) # 加载模型 # 方式一加载整个模型 model torch.load(model.pth) # 方式二先创建模型结构再加载参数 model SimpleNN() # 需要先定义同样的模型结构 model.load_state_dict(torch.load(model_weights.pth)) model.eval() # 记得设置为评估模式6. 进阶技巧让开发更高效掌握了基础用法后我再分享几个让开发更高效的小技巧6.1 使用Docker Compose管理多个服务如果你需要同时运行多个容器比如一个跑模型训练一个跑API服务可以用Docker Compose。创建一个docker-compose.yml文件version: 3.8 services: pytorch-jupyter: image: pytorch/pytorch:2.5.0-cuda12.4-cudnn9-runtime runtime: nvidia ports: - 8888:8888 volumes: - ./workspace:/workspace command: jupyter notebook --ip0.0.0.0 --allow-root --NotebookApp.token pytorch-api: build: ./api # 假设你有个API服务的Dockerfile ports: - 5000:5000 depends_on: - pytorch-jupyter然后运行docker-compose up6.2 配置Jupyter扩展让Jupyter更好用的一些扩展# 在容器内安装 pip install jupyter_contrib_nbextensions jupyter contrib nbextension install --user # 启用有用的扩展 jupyter nbextension enable code_prettify/code_prettify jupyter nbextension enable execute_time/ExecuteTime jupyter nbextension enable toggle_all_line_numbers/main6.3 使用TensorBoard可视化训练过程PyTorch支持TensorBoard可以可视化训练过程from torch.utils.tensorboard import SummaryWriter # 创建writer writer SummaryWriter(runs/experiment_1) # 在训练循环中记录 for epoch in range(num_epochs): # ... 训练代码 writer.add_scalar(Loss/train, loss.item(), epoch) writer.add_scalar(Accuracy/train, accuracy, epoch) # 在终端启动TensorBoard # tensorboard --logdirruns6.4 利用torch.compile加速PyTorch 2.x新特性PyTorch 2.5的torch.compile可以显著加速模型import torch # 普通的模型 model SimpleNN().cuda() # 编译后的模型第一次运行会慢之后会快 compiled_model torch.compile(model) # 使用编译后的模型 output compiled_model(input_data)7. 总结走完这个完整的教程你现在应该已经成功安装并运行了PyTorch 2.5镜像- 不用再折腾环境配置了验证了GPU加速可用- 看到GPU比CPU快了多少倍实际运行了三个例子- 从简单计算到完整训练流程都体验了学会了解决常见问题- 遇到问题知道怎么查怎么解决了了解了一些进阶技巧- 可以让你的开发更高效这个镜像最大的价值就是帮你跳过了深度学习中最枯燥、最易出错的环境配置环节。你不用再担心我的CUDA版本和PyTorch版本匹配吗torchvision怎么装不上为什么检测不到GPU所有这些问题镜像都帮你解决了。你拿到的是一个完整、一致、可复现的环境。无论是自己学习还是团队协作或者是部署到生产环境都用同一个镜像保证结果一致。我建议你保存好今天用到的命令和代码以后每次开始新项目都可以快速搭建环境。深度学习应该把时间花在思考模型、调整参数、分析结果上而不是浪费在配环境上。获取更多AI镜像想探索更多AI镜像和应用场景访问 CSDN星图镜像广场提供丰富的预置镜像覆盖大模型推理、图像生成、视频生成、模型微调等多个领域支持一键部署。
