从今天起,我将正式开启一个全新的 PyTorch 系列。不管你是刚入门深度学习的小白,还是从 TensorFlow 转过来的开发者,这个系列会用最通俗的语言、最可跑的代码,帮你真正把 PyTorch 用起来。
开篇咱们先解决两个问题:为什么越来越多的人选择 PyTorch? 以及 怎么在 10 分钟内搭好环境并跑通人生第一个神经网络?


一、为什么是 PyTorch?

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

  1. 动态计算图,调试像写 Python 一样自然
    PyTorch 的计算图是动态构建的,代码执行到哪,图就建到哪。这意味着你可以随时用 print 打出中间张量的值,用 pdb 下断点调试。跟调试普通 Python 程序一样,对新手极友好。

  2. 代码简洁,模型定义就是类
    写模型就是写一个 Python 类,继承 nn.Module,定义 __init__ 和 forward。结构清晰,没有魔法函数,看一眼就知道数据是怎么流动的。

  3. 学术界垄断级份额,工业界快速追赶
    现在顶会论文中 PyTorch 的占比已经碾压其他框架。工业界如 Tesla、OpenAI、Meta 等核心项目都是 PyTorch 实现。学会它,你读论文、复现模型、找工作的优势都会大很多。

  4. 生态强大且统一
    无论是视觉(torchvision)、文本(torchtext)、模型部署(TorchServe、ONNX 导出)还是分布式训练,官方都有非常成熟的配套库,不需要到处找第三方方案。

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


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

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

2.1 创建虚拟环境(强烈推荐)

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

没有 conda 也可以用 venv:

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

2.2 安装 PyTorch

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

CPU 版(所有电脑都可用)

pip install torch torchvision torchaudio

GPU 版(需要 NVIDIA 显卡+ CUDA 11.8 或 12.1)

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

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

2.3 验证安装

进入 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 完整代码(直接跑)

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 等),一个负责评估模式。

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

Logo

Agent 垂直技术社区,欢迎活跃、内容共建。

更多推荐