2. 环境搭建：PyTorch环境配置、CUDA与cuDNN安装、实验项目结构设计、日志与可视化工具（TensorBoard）集成

环境搭建这事，说难不难，说简单也真能卡你半天。我见过不少同学在配置环境上花了两天，结果发现是CUDA版本和PyTorch没对齐。咱们这一章就把这事捋清楚，一次配好，后面30章都不折腾。

2.1 PyTorch环境配置

我个人习惯用conda来管理环境。为什么？因为项目多了，每个依赖版本不一样，用conda隔离最省心。你想想看，一个项目要PyTorch 1.13，另一个要2.0，没有虚拟环境，迟早要炸。

创建环境的命令很简单：

conda create -n distiller python=3.9
conda activate distiller
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

这里有个坑——千万别用conda install torch。我曾经在项目里图省事用了conda源，结果装了个CPU版本的PyTorch，跑蒸馏的时候才发现GPU根本没干活。后来我学乖了，一律从PyTorch官网拿安装命令，版本号自己指定。

import torch
print(torch.__version__)        # 确认版本
print(torch.cuda.is_available()) # 必须返回True
print(torch.cuda.get_device_name(0)) # 看看是哪张卡

2.2 CUDA与cuDNN安装

CUDA和cuDNN的安装，说白了就是给PyTorch搭桥。PyTorch通过CUDA调用GPU，cuDNN负责加速卷积这些操作。蒸馏任务里特征图计算量很大，没有cuDNN加速，训练时间能翻倍。

安装步骤其实就三步：

1. 查驱动版本：在终端跑 nvidia-smi，看右上角的CUDA Version。比如显示12.1，那你就装CUDA 12.1的PyTorch版本。
2. 装CUDA Toolkit：去NVIDIA官网下载对应版本。我建议装runfile方式，别用deb，因为runfile可以指定安装路径，方便多版本共存。
3. 装cuDNN：下载对应CUDA版本的cuDNN包，解压后把文件复制到CUDA安装目录下。

conda install cudatoolkit=11.8
conda install cudnn

装完验证cuDNN是否生效：

python -c "import torch; print(torch.backends.cudnn.version())"

如果能打印出版本号，比如8906，说明cuDNN已经接上了。

2.3 实验项目结构设计

项目结构这事，我踩过不少坑。早期做蒸馏实验时，代码、日志、模型文件全混在一起，想复现一个结果得翻半天。后来我总结了一套结构，用了三年，团队里也一直在用。

直接看结构：

knowledge_distillation/
├── configs/              # 配置文件（YAML格式）
│   ├── student.yaml
│   └── teacher.yaml
├── data/                 # 数据集存放
│   ├── cifar100/
│   └── imagenet_subset/
├── models/               # 模型定义
│   ├── teacher/
│   │   └── resnet56.py
│   ├── student/
│   │   └── resnet20.py
│   └── __init__.py
├── losses/               # 蒸馏损失函数
│   ├── kd_loss.py
│   ├── fitnet_loss.py
│   └── attention_loss.py
├── trainers/             # 训练逻辑
│   ├── base_trainer.py
│   └── distillation_trainer.py
├── utils/                # 工具函数
│   ├── logger.py
│   └── metrics.py
├── logs/                 # 训练日志
│   └── exp_001/
├── checkpoints/          # 模型权重
│   └── best_model.pth
├── main.py               # 入口
└── requirements.txt      # 依赖清单

为什么要这么分？我举个例子。你在调蒸馏温度T这个超参数时，可能要试T=1, 2, 4, 8。如果每次改完参数还要改代码，很容易改漏。用YAML配置，改一行就行，而且每个实验的配置都保留着，复现结果时直接看配置文件就知道当时怎么跑的。

2.4 日志与可视化工具（TensorBoard）集成

蒸馏实验里，你不仅要看loss，还要看教师和学生特征图的相似度、注意力分布的变化。这些光靠打印数字是看不出来的。TensorBoard就是干这个的。

集成方式很简单，我一般这样写：

from torch.utils.tensorboard import SummaryWriter

# 初始化
writer = SummaryWriter(log_dir='logs/exp_001')

# 记录标量
writer.add_scalar('Loss/train_kd', kd_loss.item(), epoch)
writer.add_scalar('Loss/train_ce', ce_loss.item(), epoch)

# 记录特征图（蒸馏任务常用）
writer.add_histogram('Features/teacher', teacher_feat, epoch)
writer.add_histogram('Features/student', student_feat, epoch)

# 记录模型结构
dummy_input = torch.randn(1, 3, 32, 32)
writer.add_graph(student_model, dummy_input)

# 关闭
writer.close()

启动TensorBoard也很简单：

tensorboard --logdir=logs --port=6006

然后在浏览器打开 http://localhost:6006 就能看到实时曲线了。

2.5 整体环境架构图

下面这张图把整个环境搭建的脉络串起来了。从硬件驱动到上层工具，每一层都有对应关系。

这张图从左到右看，就是环境搭建的顺序。先搞定硬件驱动，再装CUDA和cuDNN，然后装PyTorch，最后搭项目结构和工具。每一步都依赖上一步，跳不过去。

2.6 环境验证清单

配完环境后，我建议你跑一遍这个清单，确保所有组件都正常工作：