从零到一：在曙光超算上跑通你的第一个PyTorch深度学习模型（DCU/GPU双版本保姆级教程）

从零到一：在曙光超算上跑通你的第一个PyTorch深度学习模型（DCU/GPU双版本保姆级教程）

第一次接触超算平台的研究者，往往会被复杂的作业提交流程和环境配置吓退。本文将带你一步步完成从环境准备到模型训练的全过程，针对曙光超算的DCU和GPU两种计算设备提供完整操作指南。

1. 环境准备：选择适合的计算设备

曙光超算平台同时支持国产DCU和英伟达GPU两种加速设备。在开始之前，需要明确你的项目更适合哪种硬件：

对比项	DCU	GPU
软件生态	需使用定制版PyTorch	支持官方原生PyTorch
计算稳定性	可能存在波动	通常表现稳定
适用场景	国产化需求项目	需要快速验证的研发项目
安装复杂度	需配置额外环境变量	标准CUDA环境

提示：如果是首次尝试，建议从GPU版本开始，遇到问题更容易找到解决方案。

1.1 基础环境配置

无论选择哪种设备，都需要先完成这些基础步骤：

1. 登录超算平台：

   ssh [e-mail]

2. 创建独立Python环境：

   conda create -n pytorch_demo python=3.7 conda activate pytorch_demo

3. 安装基础依赖：

   pip install numpy pandas matplotlib

2. DCU专属环境配置

2.1 安装定制版PyTorch

DCU需要安装曙光提供的专用版本：

pip install /public/software/apps/DeepLearning/whl/dtk-22.04.2/torch-1.10.0a0_gitc7f69d6_dtk22.04.2-cp37-cp37m-manylinux2014_x86_64.whl pip install /public/software/apps/DeepLearning/whl/dtk-22.04.2/torchvision-0.10.0a0_dtk22.04.2_e17f5ea-cp37-cp37m-manylinux2014_x86_64.whl

2.2 环境变量配置

创建环境变量配置文件：

vi ~/pytorch_env.sh

添加以下内容：

export LD_LIBRARY_PATH=/public/software/apps/DeepLearning/PyTorch_Lib/lib:/public/software/apps/DeepLearning/PyTorch_Lib/lmdb-0.9.24-build/lib:/public/software/apps/DeepLearning/PyTorch_Lib/opencv-2.4.13.6-build/lib:/public/software/apps/DeepLearning/PyTorch_Lib/openblas-0.3.7-build/lib:$LD_LIBRARY_PATH

激活配置：

source ~/pytorch_env.sh

2.3 验证安装

申请计算节点并验证：

salloc -p dcu -N 1 --gres=dcu:2 ssh compute_node_address module switch compiler/dtk/22.04.1 python -c "import torch; print(torch.cuda.is_available(), torch.__version__)"

3. GPU环境配置

3.1 安装标准PyTorch

GPU用户可以直接安装官方版本：

conda install pytorch torchvision torchaudio cudatoolkit=11.3 -c pytorch

3.2 环境准备

申请GPU节点并加载CUDA：

salloc -p gpu -N 1 --gres=gpu:1 ssh compute_node_address module load apps/cuda/11.3

3.3 验证安装

python -c "import torch; print(torch.cuda.is_available(), torch.__version__)"

4. 作业提交与管理

4.1 准备作业脚本

创建作业提交脚本train.sh：

#!/bin/bash #SBATCH -J pytorch_demo #SBATCH -p dcu/gpu #SBATCH -N 1 #SBATCH --gres=dcu:2 # 或gpu:1 # 加载环境 module switch compiler/dtk/22.04.1 # DCU专用 # module load apps/cuda/11.3 # GPU专用 source activate pytorch_demo # 运行训练脚本 python train.py

4.2 提交与监控作业

提交作业：

sbatch train.sh

查看作业状态：

squeue -u $USER

查看实时输出：

tail -f slurm-<jobid>.out

取消作业：

scancel <jobid>

5. 实战案例：MNIST训练

5.1 准备训练脚本

创建train.py：

import torch import torch.nn as nn import torch.optim as optim from torchvision import datasets, transforms # 定义模型 class Net(nn.Module): def __init__(self): super(Net, self).__init__() self.fc = nn.Linear(784, 10) def forward(self, x): return self.fc(x.view(-1, 784)) # 数据加载 transform = transforms.Compose([transforms.ToTensor()]) train_data = datasets.MNIST('./data', train=True, download=True, transform=transform) train_loader = torch.utils.data.DataLoader(train_data, batch_size=64, shuffle=True) # 训练流程 device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = Net().to(device) optimizer = optim.SGD(model.parameters(), lr=0.01) for epoch in range(10): for data, target in train_loader: data, target = data.to(device), target.to(device) optimizer.zero_grad() output = model(data) loss = nn.CrossEntropyLoss()(output, target) loss.backward() optimizer.step() print(f'Epoch {epoch}, Loss: {loss.item():.4f}')

5.2 常见问题排查

• DCU计算结果不稳定：尝试减少batch size或调整学习率
• GPU内存不足：检查--gres=gpu:1参数是否正确
• 模块加载失败：确认module load的版本与安装版本一致

6. 性能优化技巧

1. 数据预处理优化：

   • 使用torch.utils.data.DataLoader的num_workers参数
   • 预先把数据解压到本地存储

2. 混合精度训练：

   from torch.cuda.amp import GradScaler, autocast scaler = GradScaler() with autocast(): output = model(data) loss = criterion(output, target) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()

3. 作业调度策略：

   • 使用--time=HH:MM:SS指定合理运行时间
   • 通过--mem=XXG申请足够内存

在实际项目中，我发现DCU版本对环境变量配置非常敏感，建议将关键环境变量配置写入.bashrc文件。而GPU版本通常更稳定，适合快速验证模型效果。