小白友好PyTorch-CUDA-v2.9镜像使用全攻略Jupyter和SSH详细教程你是不是刚接触深度学习被PyTorch、CUDA、环境配置这些名词搞得晕头转向或者你已经是个老手但每次在新机器上搭建环境都要折腾半天不是版本冲突就是驱动问题如果你正在寻找一个“开箱即用”的深度学习环境能让你跳过繁琐的安装步骤直接开始写代码、跑模型那么你来对地方了。今天我要介绍的PyTorch-CUDA-v2.9镜像就是为你准备的终极解决方案。这是一个预装了PyTorch 2.9、CUDA工具包和所有必要依赖的完整环境打包在一个Docker镜像里。你不需要手动安装任何东西只需要一条命令就能获得一个完全隔离、功能齐全的深度学习开发环境。更重要的是这个镜像还内置了两种最常用的开发方式Jupyter Notebook和SSH远程连接。无论你是喜欢在浏览器里交互式编程还是习惯在终端里敲命令都能找到最适合你的工作方式。这篇文章就是一份“保姆级”的使用指南。我会用最直白的语言手把手教你如何启动这个镜像如何使用Jupyter和SSH以及如何在这个环境里开始你的第一个深度学习项目。即使你之前完全没接触过Docker也能跟着步骤一步步完成。1. 为什么你需要这个镜像告别环境配置的噩梦在深入具体操作之前我们先搞清楚一个问题为什么不用传统的手动安装方式而要选择这个镜像想象一下这个场景你从GitHub上找到了一个很酷的模型代码兴冲冲地准备复现。结果第一步“安装依赖”就卡住了——PyTorch版本不对CUDA版本不匹配某个Python包死活装不上。折腾了几个小时环境还没配好热情已经消磨了一大半。PyTorch-CUDA-v2.9镜像就是为了彻底解决这个问题而生的。1.1 传统安装 vs 镜像使用天壤之别让我们用一个简单的表格对比一下对比维度传统手动安装PyTorch-CUDA-v2.9镜像安装时间几小时到几天取决于网络和问题几分钟一条命令搞定环境一致性每台电脑都可能不一样容易出“玄学”问题完全一致在任何机器上运行结果相同依赖管理容易产生包冲突需要虚拟环境隔离天然隔离每个容器都是独立环境清理卸载很难彻底清除可能残留文件影响系统一键删除容器删除后不留痕迹团队协作每人重复配置文档再详细也难免出错共享同一个镜像确保所有人环境一致GPU支持需要手动安装驱动、CUDA、cuDNN版本必须严格匹配预装且匹配好直接支持GPU加速1.2 这个镜像里有什么当你拉取这个镜像时你得到的是一个完整的、立即可用的深度学习工作站PyTorch 2.9最新的稳定版本包含所有新特性和性能优化CUDA ToolkitNVIDIA GPU加速的核心版本已经与PyTorch完美匹配cuDNN深度神经网络加速库训练速度大幅提升Python科学计算栈NumPy、Pandas、Matplotlib等常用库Jupyter Lab交互式编程环境在浏览器里写代码、看结果SSH服务远程命令行访问适合长时间运行的任务常用工具git、vim、tmux等开发工具简单来说你需要的一切都已经准备好了。你不用关心“先装A还是先装B”也不用担心“版本X和版本Y是否兼容”。这一切镜像都帮你搞定了。2. 准备工作安装Docker和NVIDIA驱动在开始使用镜像之前你需要确保你的电脑上已经安装了两个必要的东西Docker和NVIDIA显卡驱动。2.1 安装DockerDocker是运行容器的基础。如果你还没安装可以按照以下步骤在Ubuntu上安装# 更新软件包索引 sudo apt-get update # 安装必要的依赖 sudo apt-get install apt-transport-https ca-certificates curl software-properties-common # 添加Docker官方GPG密钥 curl -fsSL https://download.docker.com/linux/ubuntu/gpg | sudo apt-key add - # 添加Docker仓库 sudo add-apt-repository deb [archamd64] https://download.docker.com/linux/ubuntu $(lsb_release -cs) stable # 再次更新并安装Docker sudo apt-get update sudo apt-get install docker-ce docker-ce-cli containerd.io # 将当前用户添加到docker组避免每次用sudo sudo usermod -aG docker $USER # 需要重新登录使更改生效在Windows/Mac上安装直接去Docker官网https://www.docker.com/products/docker-desktop下载Docker Desktop安装包按照向导安装即可。Windows用户需要确保开启了WSL 2或Hyper-V。安装完成后打开终端输入以下命令验证docker --version如果显示版本号如Docker version 24.0.7说明安装成功。2.2 安装NVIDIA驱动和Docker GPU支持如果你有NVIDIA显卡并且想用GPU加速深度学习强烈推荐还需要额外步骤1. 安装NVIDIA显卡驱动# Ubuntu上可以使用ubuntu-drivers工具 sudo ubuntu-drivers autoinstall # 或者去NVIDIA官网下载对应驱动 # https://www.nvidia.com/Download/index.aspx安装后重启电脑然后验证nvidia-smi你应该能看到显卡信息类似这样----------------------------------------------------------------------------- | NVIDIA-SMI 535.161.07 Driver Version: 535.161.07 CUDA Version: 12.2 | |--------------------------------------------------------------------------- | GPU Name Persistence-M| Bus-Id Disp.A | Volatile Uncorr. ECC | | Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. | | | | MIG M. | || | 0 NVIDIA GeForce RTX 4090 Off | 00000000:01:00.0 On | Off | | 0% 38C P8 20W / 450W | 125MiB / 24564MiB | 0% Default | ---------------------------------------------------------------------------2. 安装NVIDIA Container Toolkit这是让Docker容器能使用GPU的关键# 添加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 # 安装nvidia-container-toolkit sudo apt-get update sudo apt-get install -y nvidia-container-toolkit # 重启Docker服务 sudo systemctl restart docker验证GPU在Docker中可用docker run --rm --gpus all nvidia/cuda:12.2.0-base-ubuntu20.04 nvidia-smi如果看到和主机一样的显卡信息说明配置成功。3. 获取和启动PyTorch-CUDA-v2.9镜像准备工作完成后现在我们来获取并启动镜像。3.1 拉取镜像通常镜像会托管在某个容器仓库中。假设镜像名为pytorch-cuda:v2.9拉取命令如下docker pull pytorch-cuda:v2.9这个过程会下载镜像的所有层第一次下载可能需要一些时间几个GB取决于你的网速。下载完成后你可以查看本地已有的镜像docker images应该能看到pytorch-cuda:v2.9在列表中。3.2 启动容器两种模式根据你的使用习惯可以选择不同的启动方式方式一交互式模式适合测试和临时使用docker run --gpus all -it --rm pytorch-cuda:v2.9 /bin/bash--gpus all让容器能使用所有GPU-it交互式终端模式--rm容器退出后自动删除不保存状态/bin/bash启动bash shell运行后你会直接进入容器的命令行可以开始操作了。方式二后台运行模式适合长期使用docker run --gpus all -d --name pytorch-env \ -p 8888:8888 \ -p 2222:22 \ -v $(pwd):/workspace \ pytorch-cuda:v2.9-d后台运行--name pytorch-env给容器起个名字方便管理-p 8888:8888将容器的8888端口映射到主机的8888端口给Jupyter用-p 2222:22将容器的22端口映射到主机的2222端口给SSH用-v $(pwd):/workspace将当前目录挂载到容器的/workspace目录这样文件可以双向同步查看运行中的容器docker ps停止容器docker stop pytorch-env重新启动docker start pytorch-env进入已运行的容器docker exec -it pytorch-env /bin/bash4. 使用Jupyter Notebook浏览器里的交互式开发Jupyter Notebook是数据科学和机器学习领域最受欢迎的工具之一。它让你能在浏览器里写代码、运行代码、查看结果还能插入文字说明、公式和图片非常适合教学、探索和调试。4.1 启动Jupyter服务如果你按照上面的“后台运行模式”启动了容器Jupyter服务应该已经在运行了。如果没有你可以在容器内手动启动# 进入容器 docker exec -it pytorch-env /bin/bash # 启动Jupyter Lab推荐功能更强大 jupyter lab --ip0.0.0.0 --port8888 --allow-root --no-browser # 或者启动经典的Jupyter Notebook jupyter notebook --ip0.0.0.0 --port8888 --allow-root --no-browser启动后会看到类似这样的输出[I 2024-05-20 10:30:15.123 LabApp] JupyterLab extension loaded from /usr/local/lib/python3.8/dist-packages/jupyterlab [I 2024-05-20 10:30:15.123 LabApp] JupyterLab application directory is /usr/local/share/jupyter/lab [I 2024-05-20 10:30:15.128 ServerApp] jupyterlab | extension was successfully loaded. [I 2024-05-20 10:30:15.130 ServerApp] Serving notebooks from local directory: /workspace [I 2024-05-20 10:30:15.130 ServerApp] Jupyter Server 1.23.4 is running at: [I 2024-05-20 10:30:15.130 ServerApp] http://hostname:8888/lab?tokenabcdef1234567890abcdef1234567890abcdef12 [I 2024-05-20 10:30:15.130 ServerApp] or http://127.0.0.1:8888/lab?tokenabcdef1234567890abcdef1234567890abcdef12 [I 2024-05-20 10:30:15.130 ServerApp] Use Control-C to stop this server and shut down all kernels (twice to skip confirmation).注意那个tokenabcdef1234567890abcdef1234567890abcdef12你的实际token会不同等会儿登录要用。4.2 访问Jupyter界面打开你的浏览器输入http://localhost:8888或者如果Jupyter运行在远程服务器上http://服务器IP地址:8888你会看到登录页面需要输入上面看到的token输入token后点击登录就进入了Jupyter Lab界面4.3 在Jupyter中开始第一个PyTorch程序现在让我们创建一个新的Notebook写一个简单的PyTorch程序来验证环境是否正常工作。在Jupyter Lab左侧文件浏览器中点击图标创建一个新的Python 3 Notebook在第一个单元格中输入以下代码import torch # 检查PyTorch版本 print(fPyTorch版本: {torch.__version__}) # 检查CUDA是否可用 print(fCUDA是否可用: {torch.cuda.is_available()}) if torch.cuda.is_available(): # 显示GPU信息 print(fGPU数量: {torch.cuda.device_count()}) print(f当前GPU: {torch.cuda.current_device()}) print(fGPU名称: {torch.cuda.get_device_name(0)}) # 显示GPU内存信息 print(fGPU总内存: {torch.cuda.get_device_properties(0).total_memory / 1e9:.2f} GB) print(fGPU已用内存: {torch.cuda.memory_allocated(0) / 1e9:.2f} GB) print(fGPU保留内存: {torch.cuda.memory_reserved(0) / 1e9:.2f} GB)按ShiftEnter运行这个单元格你应该看到类似这样的输出PyTorch版本: 2.9.0 CUDA是否可用: True GPU数量: 1 当前GPU: 0 GPU名称: NVIDIA GeForce RTX 4090 GPU总内存: 24.00 GB GPU已用内存: 0.00 GB GPU保留内存: 0.00 GB恭喜你的PyTorchCUDA环境已经正常工作了。4.4 一个完整的深度学习示例让我们再试一个更有趣的例子——用GPU训练一个简单的神经网络import torch import torch.nn as nn import torch.optim as optim import torch.nn.functional as F from torch.utils.data import DataLoader, TensorDataset import matplotlib.pyplot as plt import numpy as np # 设置随机种子保证可重复性 torch.manual_seed(42) if torch.cuda.is_available(): torch.cuda.manual_seed(42) # 创建一些模拟数据 n_samples 1000 n_features 20 n_classes 3 # 生成数据 X torch.randn(n_samples, n_features) # 创建一个简单的分类任务 y torch.randint(0, n_classes, (n_samples,)) # 创建数据集和数据加载器 dataset TensorDataset(X, y) dataloader DataLoader(dataset, batch_size32, shuffleTrue) # 定义一个简单的神经网络 class SimpleNN(nn.Module): def __init__(self, input_size, hidden_size, num_classes): super(SimpleNN, self).__init__() self.fc1 nn.Linear(input_size, hidden_size) self.fc2 nn.Linear(hidden_size, hidden_size // 2) self.fc3 nn.Linear(hidden_size // 2, num_classes) self.dropout nn.Dropout(0.3) def forward(self, x): x F.relu(self.fc1(x)) x self.dropout(x) x F.relu(self.fc2(x)) x self.dropout(x) x self.fc3(x) return x # 创建模型如果有GPU就放到GPU上 device torch.device(cuda if torch.cuda.is_available() else cpu) model SimpleNN(n_features, 64, n_classes).to(device) # 定义损失函数和优化器 criterion nn.CrossEntropyLoss() optimizer optim.Adam(model.parameters(), lr0.001) # 训练模型 n_epochs 10 train_losses [] train_accuracies [] model.train() for epoch in range(n_epochs): epoch_loss 0.0 correct 0 total 0 for batch_idx, (data, target) in enumerate(dataloader): data, target data.to(device), target.to(device) # 前向传播 optimizer.zero_grad() output model(data) loss criterion(output, target) # 反向传播 loss.backward() optimizer.step() # 统计 epoch_loss loss.item() _, predicted output.max(1) total target.size(0) correct predicted.eq(target).sum().item() # 计算本epoch的指标 avg_loss epoch_loss / len(dataloader) accuracy 100. * correct / total train_losses.append(avg_loss) train_accuracies.append(accuracy) print(fEpoch [{epoch1}/{n_epochs}], Loss: {avg_loss:.4f}, Accuracy: {accuracy:.2f}%) print(训练完成) # 可视化训练过程 fig, (ax1, ax2) plt.subplots(1, 2, figsize(12, 4)) ax1.plot(train_losses, label训练损失) ax1.set_xlabel(Epoch) ax1.set_ylabel(Loss) ax1.set_title(训练损失曲线) ax1.legend() ax1.grid(True) ax2.plot(train_accuracies, label训练准确率, colororange) ax2.set_xlabel(Epoch) ax2.set_ylabel(Accuracy (%)) ax2.set_title(训练准确率曲线) ax2.legend() ax2.grid(True) plt.tight_layout() plt.show()运行这个代码你会看到模型在训练损失在下降准确率在上升最后还会显示两张图表。这一切都在GPU上运行速度比CPU快得多。4.5 Jupyter使用小技巧快捷键ShiftEnter运行当前单元格并跳到下一个CtrlEnter运行当前单元格并留在当前单元格Esc进入命令模式A在上面插入单元格B在下面插入单元格DD删除单元格M将单元格转为Markdown写文档用Y将单元格转为代码魔法命令%time测量单行代码的执行时间%timeit多次运行取平均时间%matplotlib inline在Notebook中显示图表%%writefile filename.py将单元格内容保存为文件文件管理在Jupyter Lab中可以直接上传、下载、重命名、移动文件支持打开多种文件类型.py、.ipynb、.txt、.csv、.jpg等5. 使用SSH连接终端里的高效开发虽然Jupyter很强大但有些时候你可能更喜欢在终端里工作。比如运行需要长时间的训练任务使用tmux或screen管理多个会话进行文件系统操作使用vim等终端编辑器这时候SSH就派上用场了。5.1 配置SSH访问首先确保你的容器已经映射了SSH端口。回顾一下我们的启动命令docker run --gpus all -d --name pytorch-env \ -p 8888:8888 \ -p 2222:22 \ # 这里将容器的22端口映射到主机的2222端口 -v $(pwd):/workspace \ pytorch-cuda:v2.9容器内部已经安装了SSH服务。我们需要设置一个密码或者更好的方式——使用SSH密钥。方法一设置密码登录进入容器docker exec -it pytorch-env /bin/bash设置root密码passwd输入并确认你的密码。确保SSH服务正在运行service ssh start方法二使用SSH密钥更安全在主机上生成SSH密钥如果还没有ssh-keygen -t rsa -b 4096将公钥复制到容器中# 先进入容器 docker exec -it pytorch-env /bin/bash # 创建.ssh目录 mkdir -p /root/.ssh # 退出容器在主机上复制公钥 docker cp ~/.ssh/id_rsa.pub pytorch-env:/root/.ssh/authorized_keys # 回到容器设置权限 docker exec -it pytorch-env /bin/bash chmod 700 /root/.ssh chmod 600 /root/.ssh/authorized_keys5.2 连接容器现在可以从主机SSH连接到容器了ssh rootlocalhost -p 2222如果设置了密码会提示输入密码如果使用了密钥应该直接连接成功。连接成功后你会看到容器的命令行提示符root容器ID:~#5.3 在SSH会话中工作一旦连接成功你就可以像在本地终端一样工作了查看GPU状态nvidia-smi运行Python脚本python train.py使用tmux管理会话强烈推荐# 启动一个新的tmux会话 tmux new -s training # 在tmux中运行长时间任务 python long_training.py # 按Ctrlb然后按d分离会话任务在后台继续运行 # 重新连接会话 tmux attach -t training # 查看所有会话 tmux ls # 结束会话 tmux kill-session -t training文件操作# 查看目录 ls -la # 因为我们在启动时挂载了当前目录到/workspace cd /workspace ls # 这里应该能看到主机当前目录的文件 # 编辑文件 vim train.py # 或者用nano nano train.py5.4 一个实际的训练脚本示例让我们创建一个简单的训练脚本然后在SSH会话中用tmux运行它创建train.py#!/usr/bin/env python3 一个简单的PyTorch训练脚本演示如何在SSH会话中运行长时间任务 import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader, TensorDataset import time import os # 设置设备 device torch.device(cuda if torch.cuda.is_available() else cpu) print(f使用设备: {device}) # 创建模拟数据 n_samples 10000 n_features 100 X torch.randn(n_samples, n_features) y torch.randn(n_samples, 1) dataset TensorDataset(X, y) dataloader DataLoader(dataset, batch_size64, shuffleTrue) # 定义一个更复杂的模型 class ComplexModel(nn.Module): def __init__(self, input_size): super(ComplexModel, self).__init__() self.network nn.Sequential( nn.Linear(input_size, 512), nn.ReLU(), nn.Dropout(0.2), nn.Linear(512, 256), nn.ReLU(), nn.Dropout(0.2), nn.Linear(256, 128), nn.ReLU(), nn.Linear(128, 64), nn.ReLU(), nn.Linear(64, 1) ) def forward(self, x): return self.network(x) model ComplexModel(n_features).to(device) criterion nn.MSELoss() optimizer optim.Adam(model.parameters(), lr0.001) # 训练循环 n_epochs 100 print(f开始训练共{n_epochs}个epoch...) for epoch in range(n_epochs): epoch_loss 0.0 start_time time.time() for batch_idx, (data, target) in enumerate(dataloader): data, target data.to(device), target.to(device) optimizer.zero_grad() output model(data) loss criterion(output, target) loss.backward() optimizer.step() epoch_loss loss.item() avg_loss epoch_loss / len(dataloader) epoch_time time.time() - start_time # 每10个epoch保存一次模型 if (epoch 1) % 10 0: torch.save({ epoch: epoch, model_state_dict: model.state_dict(), optimizer_state_dict: optimizer.state_dict(), loss: avg_loss, }, fmodel_checkpoint_epoch_{epoch1}.pth) print(f模型已保存到 model_checkpoint_epoch_{epoch1}.pth) print(fEpoch [{epoch1:3d}/{n_epochs}], Loss: {avg_loss:.6f}, Time: {epoch_time:.2f}s) print(训练完成)在SSH中运行# 启动tmux会话 tmux new -s long_training # 运行训练脚本 python train.py # 按Ctrlb然后按d分离会话 # 现在可以关闭SSH连接训练会在后台继续 # 稍后重新连接查看进度 tmux attach -t long_training5.5 SSH连接示意图通过SSH你可以在本地终端连接到远程容器运行长时间任务而不担心网络中断同时管理多个训练任务方便地进行文件传输和编辑6. 实际工作流结合Jupyter和SSH的最佳实践在实际工作中我通常结合使用Jupyter和SSH发挥各自的优势6.1 典型工作流程探索阶段用Jupyter在Jupyter中快速尝试新想法可视化数据和模型结构调试代码逐步执行查看中间结果开发阶段用Jupyter或SSH在Jupyter中完善代码在SSH中用vim编辑脚本写单元测试确保代码质量训练阶段用SSHtmux在SSH会话中用tmux启动长时间训练定期保存检查点监控GPU使用情况和训练进度分析阶段用Jupyter加载训练好的模型可视化训练曲线分析模型性能制作报告6.2 一个完整的项目示例假设我们要做一个图像分类项目这是我会采取的步骤步骤1在Jupyter中探索数据# data_exploration.ipynb import torch from torchvision import datasets, transforms import matplotlib.pyplot as plt import numpy as np # 加载数据集 transform transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,)) ]) train_dataset datasets.MNIST(./data, trainTrue, downloadTrue, transformtransform) test_dataset datasets.MNIST(./data, trainFalse, transformtransform) # 查看数据 print(f训练集大小: {len(train_dataset)}) print(f测试集大小: {len(test_dataset)}) # 可视化一些样本 fig, axes plt.subplots(2, 5, figsize(12, 5)) for i in range(10): img, label train_dataset[i] ax axes[i//5, i%5] ax.imshow(img.squeeze(), cmapgray) ax.set_title(f标签: {label}) ax.axis(off) plt.show()步骤2在SSH中创建训练脚本# 在SSH中创建项目目录 mkdir mnist_project cd mnist_project # 用vim创建train.py vim train.pytrain.py内容import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader from torchvision import datasets, transforms import time import argparse def train(model, device, train_loader, optimizer, criterion, epoch): model.train() train_loss 0 correct 0 total 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() train_loss loss.item() _, predicted output.max(1) total target.size(0) correct predicted.eq(target).sum().item() if batch_idx % 100 0: print(fTrain Epoch: {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)} f({100. * batch_idx / len(train_loader):.0f}%)]\tLoss: {loss.item():.6f}) accuracy 100. * correct / total avg_loss train_loss / len(train_loader) return avg_loss, accuracy def test(model, device, test_loader, criterion): model.eval() test_loss 0 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) test_loss criterion(output, target).item() _, predicted output.max(1) total target.size(0) correct predicted.eq(target).sum().item() accuracy 100. * correct / total avg_loss test_loss / len(test_loader) return avg_loss, accuracy def main(): parser argparse.ArgumentParser(descriptionMNIST训练脚本) parser.add_argument(--batch-size, typeint, default64, help批大小) parser.add_argument(--epochs, typeint, default10, help训练轮数) parser.add_argument(--lr, typefloat, default0.01, help学习率) parser.add_argument(--save-dir, typestr, default./checkpoints, help模型保存目录) args parser.parse_args() # 设备设置 device torch.device(cuda if torch.cuda.is_available() else cpu) print(f使用设备: {device}) # 数据加载 transform transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)) ]) train_dataset datasets.MNIST(./data, trainTrue, downloadTrue, transformtransform) test_dataset datasets.MNIST(./data, trainFalse, transformtransform) train_loader DataLoader(train_dataset, batch_sizeargs.batch_size, shuffleTrue) test_loader DataLoader(test_dataset, batch_sizeargs.batch_size, shuffleFalse) # 模型定义 class Net(nn.Module): def __init__(self): super(Net, self).__init__() self.conv1 nn.Conv2d(1, 32, 3, 1) self.conv2 nn.Conv2d(32, 64, 3, 1) self.dropout1 nn.Dropout2d(0.25) self.dropout2 nn.Dropout2d(0.5) self.fc1 nn.Linear(9216, 128) self.fc2 nn.Linear(128, 10) def forward(self, x): x self.conv1(x) x nn.functional.relu(x) x self.conv2(x) x nn.functional.relu(x) x nn.functional.max_pool2d(x, 2) x self.dropout1(x) x torch.flatten(x, 1) x self.fc1(x) x nn.functional.relu(x) x self.dropout2(x) x self.fc2(x) return nn.functional.log_softmax(x, dim1) model Net().to(device) optimizer optim.Adam(model.parameters(), lrargs.lr) criterion nn.CrossEntropyLoss() # 训练循环 for epoch in range(1, args.epochs 1): start_time time.time() train_loss, train_acc train(model, device, train_loader, optimizer, criterion, epoch) test_loss, test_acc test(model, device, test_loader, criterion) epoch_time time.time() - start_time print(f\nEpoch {epoch}:) print(f 训练损失: {train_loss:.4f}, 训练准确率: {train_acc:.2f}%) print(f 测试损失: {test_loss:.4f}, 测试准确率: {test_acc:.2f}%) print(f 耗时: {epoch_time:.2f}秒) # 保存模型 if epoch % 5 0: torch.save({ epoch: epoch, model_state_dict: model.state_dict(), optimizer_state_dict: optimizer.state_dict(), train_acc: train_acc, test_acc: test_acc, }, f{args.save_dir}/model_epoch_{epoch}.pth) print(f模型已保存到 {args.save_dir}/model_epoch_{epoch}.pth) print(训练完成) if __name__ __main__: main()步骤3在SSH中用tmux启动训练# 创建保存目录 mkdir checkpoints # 启动tmux会话 tmux new -s mnist_training # 运行训练 python train.py --batch-size 128 --epochs 20 --lr 0.001 # 按Ctrlb然后按d分离会话步骤4在Jupyter中分析结果# analysis.ipynb import torch import matplotlib.pyplot as plt from torchvision import datasets, transforms # 加载训练好的模型 checkpoint torch.load(checkpoints/model_epoch_20.pth) # 可视化训练过程假设我们记录了历史 # 这里可以绘制损失曲线、准确率曲线等 # 测试模型性能 def visualize_predictions(model, test_loader, device, num_images10): model.eval() data_iter iter(test_loader) images, labels next(data_iter) fig, axes plt.subplots(2, 5, figsize(15, 6)) images, labels images.to(device), labels.to(device) with torch.no_grad(): outputs model(images) _, predicted outputs.max(1) for i in range(num_images): ax axes[i//5, i%5] ax.imshow(images[i].cpu().squeeze(), cmapgray) ax.set_title(f真值: {labels[i].item()}, 预测: {predicted[i].item()}) ax.axis(off) plt.tight_layout() plt.show() # 使用这个函数可视化预测结果7. 常见问题与解决方案在使用过程中你可能会遇到一些问题。这里列出了一些常见问题及其解决方法7.1 Docker相关问题问题1docker命令需要sudo解决将用户添加到docker组 sudo usermod -aG docker $USER 然后重新登录问题2端口被占用解决更改端口映射比如把-p 8888:8888改为-p 8889:8888问题3容器启动失败解决检查Docker服务是否运行 sudo systemctl status docker 如果没运行sudo systemctl start docker7.2 GPU相关问题问题1nvidia-smi在容器中不可用解决确保启动时加了--gpus all参数 docker run --gpus all ...问题2CUDA out of memory解决减小批大小或者使用梯度累积 # 在代码中 accumulation_steps 4 for i, (data, target) in enumerate(train_loader): loss criterion(model(data), target) loss loss / accumulation_steps loss.backward() if (i1) % accumulation_steps 0: optimizer.step() optimizer.zero_grad()问题3PyTorch找不到CUDA解决检查PyTorch和CUDA版本是否匹配 import torch print(torch.__version__) print(torch.version.cuda) # 应该显示CUDA版本7.3 Jupyter相关问题问题1无法访问Jupyter解决 1. 检查容器是否运行docker ps 2. 检查端口映射docker port 容器名 8888 3. 检查防火墙设置 4. 尝试用http://localhost:8888 或 http://127.0.0.1:8888问题2忘记token解决查看容器日志 docker logs 容器名 在输出中查找token问题3内核死掉解决 1. 检查内存是否不足 2. 重启内核Kernel - Restart 3. 增加资源docker run时增加内存限制7.4 SSH相关问题问题1连接被拒绝解决 1. 检查SSH服务是否运行docker exec 容器名 service ssh status 2. 检查端口映射确保-p 2222:22 3. 检查防火墙问题2密码登录失败解决 1. 在容器内重置密码passwd 2. 或使用SSH密钥问题3连接后立即断开解决检查容器资源是否充足可能是内存不足7.5 性能优化建议数据加载优化# 使用多进程数据加载 DataLoader(..., num_workers4, pin_memoryTrue) # 使用非阻塞传输 data data.to(device, non_blockingTrue)混合精度训练减少显存加快速度from torch.cuda.amp import autocast, GradScaler scaler GradScaler() for data, target in train_loader: optimizer.zero_grad() with autocast(): output model(data) loss criterion(output, target) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()梯度累积模拟大batchaccumulation_steps 4 for i, (data, target) in enumerate(train_loader): loss criterion(model(data), target) loss loss / accumulation_steps loss.backward() if (i1) % accumulation_steps 0: optimizer.step() optimizer.zero_grad()8. 总结为什么这个镜像能改变你的工作方式通过这篇详细的教程你应该已经掌握了PyTorch-CUDA-v2.9镜像的完整使用方法。让我们回顾一下这个方案的核心优势8.1 一站式解决方案这个镜像提供了一个完整的、立即可用的深度学习环境预配置好的软件栈PyTorch、CUDA、cuDNN等全部就绪两种开发方式Jupyter适合交互探索SSH适合长时间任务GPU直通支持无需额外配置就能使用GPU加速环境隔离每个项目可以有独立的环境互不干扰8.2 显著提升效率与传统方式相比使用这个镜像可以节省数小时甚至数天的环境配置时间避免在我机器上能运行的问题确保环境一致性快速切换项目每个项目一个容器互不影响轻松分享环境只需分享Dockerfile或镜像名称8.3 适合各种场景无论你是初学者想快速开始学习PyTorch不被环境问题困扰研究者需要可复现的实验环境工程师要在多台机器上部署相同的环境教师为学生提供统一的教学环境这个镜像都能满足你的需求。8.4 下一步建议现在你已经掌握了基本用法可以进一步探索定制镜像基于这个镜像创建自己的版本添加常用工具和库Docker Compose使用docker-compose管理多个服务版本控制将Dockerfile加入Git跟踪环境变化CI/CD集成在自动化流水线中使用这个镜像云部署将容器部署到云服务器随时随地访问最重要的是现在就开始使用。选择一个你感兴趣的项目用这个镜像搭建环境体验一下开箱即用的畅快感。你会发现原来深度学习环境配置可以这么简单你可以把更多时间花在真正重要的事情上——思考和创造。记住工具的目的是提高效率而不是制造障碍。PyTorch-CUDA-v2.9镜像就是这样一个工具它帮你扫清环境配置的障碍让你能专注于模型和算法本身。祝你使用愉快编码顺利获取更多AI镜像想探索更多AI镜像和应用场景访问 CSDN星图镜像广场提供丰富的预置镜像覆盖大模型推理、图像生成、视频生成、模型微调等多个领域支持一键部署。