PyTorch入门：10分钟搭建首个神经网络

从今天起，我将正式开启一个全新的 PyTorch 系列 。不管你是刚入门深度学习的小白，还是从 TensorFlow 转过来的开发者，这个系列会用最通俗的语言、最可跑的代码，帮你真正把 PyTorch 用起来。

开篇咱们先解决两个问题：为什么越来越多的人选择 PyTorch？ 以及 怎么在 10 分钟内搭好环境并跑通人生第一个神经网络？

一、为什么是 PyTorch？

深度学习框架很多，但如果你现在才开始系统学习，我建议直接从 PyTorch 入门。原因很直接：

动态计算图，调试像写 Python 一样自然

PyTorch 的计算图是动态构建的，代码执行到哪，图就建到哪。这意味着你可以随时用 print 打出中间张量的值，用 pdb 下断点调试。跟调试普通 Python 程序一样，对新手极友好。
代码简洁，模型定义就是类

写模型就是写一个 Python 类，继承 nn.Module，定义 __init__ 和 forward。结构清晰，没有魔法函数，看一眼就知道数据是怎么流动的。
学术界垄断级份额，工业界快速追赶

现在顶会论文中 PyTorch 的占比已经碾压其他框架。工业界如 Tesla、OpenAI、Meta 等核心项目都是 PyTorch 实现。学会它，你读论文、复现模型、找工作的优势都会大很多。
生态强大且统一

无论是视觉（torchvision）、文本（torchtext）、模型部署（TorchServe、ONNX 导出）还是分布式训练，官方都有非常成熟的配套库，不需要到处找第三方方案。

一句话总结：PyTorch 让你用最少的心智负担，完成最前沿的深度学习任务。

二、环境搭建：10 分钟拿下 PyTorch

不管你有 GPU 还是只有 CPU，下面的步骤都能在 10 分钟内搞定。我建议使用 conda 或 pip 安装，任选一种即可。

2.1 创建虚拟环境（强烈推荐）

bash 复制代码

# 使用 conda 创建新环境，指定 Python 版本
conda create -n pytorch_env python=3.9 -y
conda activate pytorch_env

没有 conda 也可以用 venv：

bash 复制代码

python -m venv pytorch_env
source pytorch_env/bin/activate   # Linux/Mac
# pytorch_env\Scripts\activate    # Windows

2.2 安装 PyTorch

去 PyTorch 官网选择你的配置，会自动生成安装命令。这里给出最常用的两种：

CPU 版（所有电脑都可用）

bash 复制代码

pip install torch torchvision torchaudio

GPU 版（需要 NVIDIA 显卡+ CUDA 11.8 或 12.1）

bash 复制代码

# CUDA 11.8 示例
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

如果你不确定自己 CUDA 版本，先在命令行执行 nvidia-smi 查看右上角 CUDA Version。

2.3 验证安装

进入 Python 环境，逐行执行下面的代码：

python 复制代码

import torch

print("PyTorch 版本:", torch.__version__)
print("CUDA 是否可用:", torch.cuda.is_available())
print("CUDA 版本:", torch.version.cuda)
print("GPU 设备数量:", torch.cuda.device_count())

# 做一个简单的张量运算
x = torch.rand(3, 3)
print("随机张量:\n", x)

如果 CUDA 是否可用 为 True，恭喜，你的 GPU 环境已完美就绪。如果是 False，说明你装的是 CPU 版或驱动有问题，但完全不影响接下来的学习和实验，CPU 跑 MNIST 这种小任务绰绰有余。

三、10 分钟跑通第一个网络：MNIST 手写数字识别

下面我带你用 PyTorch 训练一个识别手写数字的网络。整个流程不超过 50 行代码，复制即可运行。

3.1 完整代码（直接跑）

python 复制代码

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader

# 1. 设置设备（有 GPU 就用 GPU）
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("使用设备:", device)

# 2. 超参数
BATCH_SIZE = 64
EPOCHS = 3
LEARNING_RATE = 0.001

# 3. 数据加载与预处理
transform = transforms.Compose([
    transforms.ToTensor(),                     # 转成张量并归一化到 [0,1]
    transforms.Normalize((0.1307,), (0.3081,)) # MNIST 标准均值和标准差
])

train_dataset = torchvision.datasets.MNIST(
    root='./data', train=True, download=True, transform=transform
)
test_dataset = torchvision.datasets.MNIST(
    root='./data', train=False, download=True, transform=transform
)

train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False)

# 4. 定义网络结构（一个简单的全连接网络）
class SimpleNet(nn.Module):
    def __init__(self):
        super(SimpleNet, self).__init__()
        self.fc1 = nn.Linear(28*28, 128)   # 输入 784 个像素
        self.fc2 = nn.Linear(128, 64)
        self.fc3 = nn.Linear(64, 10)       # 输出 10 个类别(0-9)

    def forward(self, x):
        x = x.view(-1, 28*28)             # 把图像展平成一维
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)                   # 最后一层不接激活函数，因为 CrossEntropyLoss 会自带 softmax
        return x

model = SimpleNet().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)

# 5. 训练函数
def train():
    model.train()
    total_loss = 0
    for images, labels in train_loader:
        images, labels = images.to(device), labels.to(device)

        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        total_loss += loss.item()
    return total_loss / len(train_loader)

# 6. 测试函数
def test():
    model.eval()
    correct = 0
    with torch.no_grad():
        for images, labels in test_loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            _, predicted = torch.max(outputs, 1)
            correct += (predicted == labels).sum().item()
    accuracy = 100.0 * correct / len(test_dataset)
    return accuracy

# 7. 开始训练
for epoch in range(EPOCHS):
    loss_avg = train()
    acc = test()
    print(f"Epoch [{epoch+1}/{EPOCHS}] - Loss: {loss_avg:.4f}, Test Accuracy: {acc:.2f}%")

3.2 运行这个网络你会看到什么？

3 个 epoch 后，测试准确率就能达到 96%--97% 左右。一个这么简单的三层全连接网络，几秒钟就能跑出这个结果，这就是 PyTorch 的威力。

3.3 代码关键点速览

torch.device：自动选择 CPU 或 GPU，保证代码无缝迁移。
DataLoader：自动帮你分批次、打乱数据，是数据加载的标配。
nn.Module ：所有模型的基类，你只需要填 __init__ 和 forward 两个地方。
optimizer.zero_grad()：每次反向传播前一定要清空梯度，不然会累积。
loss.backward()：自动求导，链式法则全给你算好了。
model.train() / model.eval()：一个负责训练模式（开启 Dropout 等），一个负责评估模式。

如果这篇内容对你有帮助，点个赞 + 收藏 + 关注，这样更新时就能第一时间收到推送。你的支持是我持续写下去的最大动力。