📚 课程目录第1章环境搭建与基线模型准备

1. 环境搭建与基线模型准备

任务目标

搭建同时支持GPU（CUDA）与NPU（Ascend）的PyTorch深度学习环境，下载ResNet-50预训练权重作为教师模型，并验证模型在前向推理中可正常输出指定尺寸的张量。本案例为后续知识蒸馏实验提供统一的运行基座。

环境依赖与安装

本案例需要以下基础环境，安装命令在Ubuntu 20.04/22.04上验证通过：

Python 3.8+，CUDA 11.6+（GPU环境）
PyTorch 1.13+（GPU版）与 torchvision 0.14+
Ascend Toolkit 6.0+ 与 torch_npu 1.11+（NPU环境，可选）

# GPU环境安装（以CUDA 11.8为例）
pip install torch==2.0.1 torchvision==0.15.1 --index-url https://download.pytorch.org/whl/cu118

# NPU环境安装（需先安装CANN Toolkit，然后执行）
pip install torch==1.11.0 torchvision==0.12.0
pip install torch_npu==1.11.0.post5

关键完整代码

以下代码实现环境检测、教师模型加载与单次前向验证。代码同时兼容CPU/GPU/NPU，自动选择可用设备。

import torch
import torchvision.models as models
import torch_npu  # 仅在NPU环境需要，GPU环境可注释

def setup_device():
    """自动选择可用设备：NPU > GPU > CPU"""
    if torch.npu.is_available():
        return torch.device('npu:0')
    elif torch.cuda.is_available():
        return torch.device('cuda:0')
    return torch.device('cpu')

def load_teacher_model(device):
    """加载ResNet-50预训练教师模型并设为评估模式"""
    model = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V1)
    model = model.to(device)
    model.eval()
    return model

def verify_forward(model, device):
    """构造随机输入验证前向传播输出形状"""
    dummy_input = torch.randn(1, 3, 224, 224).to(device)
    with torch.no_grad():
        output = model(dummy_input)
    return output

if __name__ == '__main__':
    device = setup_device()
    print(f'使用设备: {device}')
    
    teacher = load_teacher_model(device)
    print(f'教师模型参数量: {sum(p.numel() for p in teacher.parameters()):,}')
    
    out = verify_forward(teacher, device)
    print(f'输出张量形状: {list(out.shape)}')
    print(f'输出数值范围: [{out.min().item():.3f}, {out.max().item():.3f}]')

运行与验证

将上述代码保存为 setup_baseline.py，在项目根目录执行：

python setup_baseline.py

预期输出（首次运行会自动下载约100MB权重文件）：

使用设备: cuda:0  （或 npu:0 / cpu）
教师模型参数量: 25,557,032
输出张量形状: [1, 1000]
输出数值范围: [-2.118, 6.784]

通过条件：程序无报错退出，输出张量形状为 [1, 1000]（ImageNet 1000类），且设备类型与当前硬件匹配。若使用NPU，torch.npu.is_available() 必须返回 True，否则检查CANN环境变量 source /usr/local/Ascend/ascend-toolkit/set_env.sh。

常见问题

GPU显存不足：将 batch_size 从1调小或使用 torch.cuda.empty_cache() 释放缓存。
NPU算子不支持：ResNet-50标准结构已适配，若自定义结构需检查 torch_npu 算子映射表。
权重下载超时：手动下载 resnet50-0676ba61.pth 并放置于 ~/.cache/torch/hub/checkpoints/ 目录。