第二章 环境搭建与工具链准备

说实话，每次带新人入门，我最怕的就是环境搭建这一步。明明代码逻辑没问题，结果卡在「装不上」「跑不通」这种基础问题上，特别打击信心。我自己刚入行那会儿，也曾在配置环境上折腾了整整两天——后来发现只是Python版本不对。

所以这一章，咱们把基础打牢。你跟着我的步骤来，基本不会出问题。

2.1 Python虚拟环境配置

为什么要用虚拟环境？说白了，就是隔离依赖。你同时做两个项目，一个用PyTorch 1.13，另一个用PyTorch 2.0，如果没有虚拟环境，这两个版本会打架，你想想看那得多头疼。

我个人习惯用 conda 来管理环境，当然 venv 也行。这里我推荐 conda，因为后面装CUDA相关的包会方便很多。

嗯，这里要注意一点：Python版本别选太新。我建议3.8到3.10之间，太新的版本有些算子库还没适配好。我曾经在Python 3.11上踩过坑，某个科学计算库死活装不上，最后发现是版本不兼容。

2.2 PyTorch/TensorFlow框架安装

这两个框架，咱们课程里都会用到。PyTorch做算子精度对齐更方便，TensorFlow在融合后验证上有它的优势。我的建议是：两个都装上，但以PyTorch为主。

安装命令其实很简单，但有个坑——一定要先确认你的CUDA版本。

# 查看CUDA版本
nvidia-smi

# 根据CUDA版本安装（以CUDA 11.8为例）
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

# TensorFlow安装
pip install tensorflow==2.13.0

装完之后怎么验证？我一般跑个小脚本：

import torch
print(torch.__version__)
print(torch.cuda.is_available())  # 应该返回 True

import tensorflow as tf
print(tf.__version__)
print(tf.config.list_physical_devices('GPU'))  # 应该能看到GPU信息

如果 cuda.is_available() 返回 False，别慌。大概率是CUDA版本和PyTorch版本没对上。去PyTorch官网查一下对应关系，重新装一遍就好。

2.3 NumPy/SciPy科学计算库准备

这两个库是算子对齐的基石。NumPy处理基础数组操作，SciPy提供更高级的数学函数。说白了，你写算子的时候，很多中间结果都要用NumPy来验证。

安装很简单：

pip install numpy scipy

但我建议装的时候指定版本，避免后续依赖冲突：

pip install numpy==1.24.3 scipy==1.10.1

import numpy as np

a = np.array([1.0000001, 2.0000002])
b = np.array([1.0, 2.0])

# 判断是否在1e-6的误差范围内
print(np.allclose(a, b, rtol=1e-6))  # True

2.4 日志与可视化工具安装

做算子对齐，可视化是必不可少的。你光看数字，很难直观感受到误差分布。Matplotlib就是干这个的。

安装命令：

pip install matplotlib

日志方面，我推荐用Python自带的 logging 模块，轻量且够用。当然你也可以用 loguru，更现代化一些。

import logging

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('operator_align.log'),
        logging.StreamHandler()
    ]
)

logger = logging.getLogger(__name__)
logger.info("算子精度对齐开始...")

可视化这块，我一般用Matplotlib画误差分布直方图和散点图。比如对比两个算子的输出差异：

import matplotlib.pyplot as plt
import numpy as np

# 模拟两组数据
output_ref = np.random.randn(1000)
output_test = output_ref + np.random.randn(1000) * 0.01

# 计算误差
error = output_test - output_ref

# 画直方图
plt.figure(figsize=(10, 6))
plt.hist(error, bins=50, alpha=0.7, color='steelblue')
plt.xlabel('误差值')
plt.ylabel('频次')
plt.title('算子输出误差分布')
plt.grid(True, alpha=0.3)
plt.show()

2.5 本章知识体系总览

为了让你对整个环境搭建有个全局认识，我画了张图：

2.6 验证环境是否就绪

所有东西装完之后，我建议你跑一个完整的验证脚本。别等到写代码时才发现某个库没装好，那太浪费时间了。

# verify_env.py
import sys
import numpy as np
import scipy
import torch
import tensorflow as tf
import matplotlib
import logging

print(f"Python版本: {sys.version}")
print(f"NumPy版本: {np.__version__}")
print(f"SciPy版本: {scipy.__version__}")
print(f"PyTorch版本: {torch.__version__}")
print(f"CUDA可用: {torch.cuda.is_available()}")
print(f"TensorFlow版本: {tf.__version__}")
print(f"Matplotlib版本: {matplotlib.__version__}")

# 简单测试
a = np.array([1.0, 2.0, 3.0])
b = np.array([1.0, 2.0, 3.0])
assert np.allclose(a, b), "NumPy基础测试失败"

x = torch.tensor([1.0, 2.0, 3.0])
y = torch.tensor([1.0, 2.0, 3.0])
assert torch.allclose(x, y), "PyTorch基础测试失败"

print("\n✅ 所有环境验证通过！")

如果这个脚本能顺利跑完，恭喜你，环境搭建成功了。如果报错，别急——看看错误信息，大概率是版本问题或者CUDA没配置好。去网上搜一下，基本都能解决。