用 PyTorch 从张量到训练一个神经网络

PyTorch 是一个开源的深度学习框架，核心特点是动态计算图：网络结构在运行时构建，因此可以在代码里直接修改模型行为，调试时也能逐行观察张量的变化。它适合两类读者：一是刚接触深度学习、想亲手把「张量 → 数据管道 → 训练循环 → 评估」这条链路走通的人；二是已经会用别的框架、想快速迁移到 PyTorch 的开发者。本文按这条链路组织，每一节都给出可直接运行的代码。如果你更想先确认环境与依赖是否就绪，可以先看 PyTorch 工具页，再回到本文动手。

准备工作

Python 环境与安装

PyTorch 通过 pip 安装，最基础的两个包是 torch 与 torchvision：

pip install torch torchvision

如果你的机器上有支持 CUDA 的 GPU，需要安装与显卡驱动匹配的 CUDA 版本。素材给出的示例命令是：

pip install torch torchvision torchaudio cudatoolkit=11.3

这里的 cudatoolkit=11.3 只是一个示例版本号，实际应安装与你本机驱动、操作系统相匹配的版本。具体可用的版本组合、安装命令生成方式以及各版本的生命周期，请以 PyTorch 官网当前信息为准，不要照抄本文的版本号。

验证安装

安装完成后，用下面这段代码确认 PyTorch 能正常导入，并判断当前环境是否可用 GPU：

import torch

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f'Using device: {device}')

这段判断逻辑在后面的 GPU 加速一节还会用到。它把设备选择收敛到一个变量里，后续创建张量时统一传入 device，代码就能在有 GPU 和没有 GPU 的机器上通用。

前置知识

阅读本文需要会写基本的 Python 函数与类，理解「输入 → 输出」的函数映射即可。神经网络部分的数学只用到矩阵乘法与梯度下降的直观概念，不要求推导。

操作步骤

第一步：创建张量

张量（tensor）是多维数组，是 PyTorch 里一切数据的基本结构。可以用 Python 列表直接构造，也可以用内置的工厂函数生成特定形状与初始值的张量：

import torch

tensor_1d = torch.tensor([1, 2, 3])
print("1D Tensor (Vector):")
print(tensor_1d)
print()

tensor_2d = torch.tensor([[1, 2], [3, 4]])
print("2D Tensor (Matrix):")
print(tensor_2d)
print()

random_tensor = torch.rand(2, 3)
print("Random Tensor (2x3):")
print(random_tensor)
print()

zeros_tensor = torch.zeros(2, 3)
print("Zeros Tensor (2x3):")
print(zeros_tensor)
print() 3)
print("Ones Tensor (2x3):")
print(ones_tensor)

几个函数的区别值得记牢：torch.tensor() 从已有数据构造，数据类型由数据推断；torch.rand() 生成指定形状的随机张量；torch.zeros() 与 torch.ones() 分别生成全 0 与全 1 张量。后三个函数只接收形状参数，不接收数据。

第二步：索引、切片与变形

拿到张量之后，最常见的三类操作是取单个元素、取子区域、改变形状。下面用一个 3×2 的张量演示：

import torch

tensor = torch.tensor([[1, 2], [3, 4], [5, 6]])

element = tensor[1, 0]
print(f"Indexed Element (Row 1, Column 0): {element}")

slice_tensor = tensor[:2, :]
print(f"Sliced Tensor (First two rows): \n{slice_tensor}")

reshaped_tensor = tensor.view(2, 3)
print(f"Reshaped Tensor (2x3): \n{reshaped_tensor}")

• 索引：tensor[1, 0] 取第 1 行第 0 列的元素。注意 PyTorch 的行列下标从 0 开始。
• 切片：tensor[:2, :] 取前两行的全部列，得到一个新的 2×2 张量。
• 变形：view(2, 3) 把 3×2 的张量重排成 2×3，元素总数不变、数值不变，只是维度组织方式变了。

变形操作要求元素总数前后一致，否则会报错。这一点在把图像数据展平送进全连接层时尤其容易踩到。

第三步：广播与矩阵乘法

广播（broadcasting）让形状不同的张量也能做逐元素运算：较小的张量会在缺失的维度上自动扩展。矩阵乘法则是神经网络里最核心的计算，用 torch.matmul 完成：

import torch

tensor_a = torch.tensor([[1, 2, 3], [4, 5, 6]])
tensor_b = torch.tensor([[10, 20, 30]])

broadcasted_result = tensor_a + tensor_b
print(f"Broadcasted Addition Result: \n{broadcasted_result}")

matrix_multiplication_result = torch.matmul(tensor_a, tensor_a.T)
print(f"Matrix Multiplication Result (tensor_a * tensor_a^T): \n{matrix_multiplication_result}")

tensor_a 的形状是 2×3，tensor_b 的形状是 1×3。相加时 tensor_b 被扩展成 2×3，两行分别加上 [10, 20, 30]。矩阵乘法里 tensor_a.T 是转置，形状变为 3×2，因此 2×3 乘 3×2 得到 2×2 的结果。理解形状如何变化，是排查神经网络维度报错的关键。

第四步：把计算放到 GPU 上

深度学习涉及大量矩阵运算，把张量搬到 GPU 上通常能显著缩短训练时间。做法是在创建张量时传入 device 参数：

import torch

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f'Using device: {device}')

tensor_size = (10000, 10000)
a = torch.randn(tensor_size, device=device)
b = torch.randn(tensor_size, device=device)
c = a + b

print("Result shape (moved to CPU for printing):", c.cpu().shape)
print("Current GPU memory usage:")
print(f"Allocated: {torch.cuda.memory_allocated(device) / (1024 ** 2):.2f} MB")
print(f"Cached: {torch.cuda.memory_reserved(device) / (1024 ** 2):.2f} MB")

这里有两个实用细节。第一，打印 GPU 上的张量前先调用 .cpu()，把结果搬回内存，避免在部分环境下打印出错。第二，torch.cuda.memory_allocated() 返回当前实际被张量占用的显存，torch.cuda.memory_reserved() 返回缓存分配器保留的显存，两者都以字节为单位，除以 1024 ** 2 换算成 MB。显存不足是训练大模型时最常见的报错来源，这两个数值是排查的起点。

如果 torch.cuda.is_available() 返回 False，上面的代码会自动回退到 CPU，逻辑仍然正确，只是速度慢。这也是把设备选择写成变量的好处。

第五步：定义神经网络类

PyTorch 的模型都继承自 torch.nn.Module，需要实现两个部分：在 __init__ 里声明各层，在 forward 里定义数据如何流经这些层。下面是一个含一个隐藏层的网络，输入 2 维、隐藏层 4 维、输出 1 维：

import torch
import torch.nn as nn

class SimpleNN(nn.Module):
    def __init__(self):
        super(SimpleNN, self).__init__()
        self.fc1 = nn.Linear(2, 4)
        self.fc2 = nn.Linear(4, 1)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

nn.Linear(2, 4) 表示输入特征数为 2、输出特征数为 4 的全连接层，内部自带权重与偏置。forward 里先过第一层，再套 ReLU 激活函数引入非线性，最后过第二层得到输出。没有激活函数的话，两层线性变换叠加等价于一层线性变换，网络就失去了拟合非线性关系的能力。

第六步：准备训练数据

为了能直观判断模型是否学到了东西，这里用 XOR（异或）逻辑门作为数据集：输入是两个二进制位，输出是它们的异或结果。XOR 是非线性可分的，正好能检验隐藏层是否起作用：

X_train = torch.tensor([[0.0, 0.0], [0.0, 1.0], [1.0, 0.0], [1.0, 1.0]])
y_train = torch.tensor([[0.0], [1.0], [1.0], [0.0]])

注意输入与标签都写成浮点数，并且标签的形状是 4×1 而不是长度 4 的一维张量，这样才能和模型输出 (4, 1) 的形状对齐。

第七步：实例化模型、损失函数与优化器

模型、损失函数、优化器是训练的三件套。这里用均方误差作为损失，用随机梯度下降（SGD）更新权重，学习率设为 0.1：

import torch.optim as optim

model = SimpleNN()
criterion = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.1)

model.parameters() 把模型里所有需要训练的张量交给优化器，优化器据此知道该更新谁。lr 是学习率，控制每次更新的步长。

第八步：编写训练循环

训练循环是 PyTorch 里最固定的骨架，顺序不能乱：前向传播算输出 → 算损失 → 清空上一步的梯度 → 反向传播算梯度 → 优化器更新参数。

for epoch in range(100):
    model.train()
    outputs = model(X_train)
    loss = criterion(outputs, y_train)

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    if (epoch + 1) % 10 == 0:
        print(f'Epoch [{epoch + 1}/100], Loss: {loss.item():.4f}')

几个容易出错的点：optimizer.zero_grad() 必须放在 loss.backward() 之前，因为 PyTorch 默认会累加梯度，不清零就会把上一轮的梯度叠加进来；loss.item() 把单元素张量取出为 Python 浮点数，方便格式化打印；model.train() 把模型切到训练模式，对含 Dropout、BatchNorm 的模型这一步是必须的。每 10 个 epoch 打印一次损失，可以观察损失是否在下降。

第九步：评估模型

训练结束后要评估模型在新数据上的表现。评估时切到 model.eval()，并用 torch.no_grad() 关闭梯度计算，既省内存也更快：

model.eval()
with torch.no_grad():
    test_data = torch.tensor([[0.0, 0.0], [0.0, 1.0], [1.0, 0.0], [1.0, 1.0]])
    predictions = model(test_data)
    print(f'Predictions: \n{predictions}')

输出是四个连续值，分别对应四组输入的预测。因为损失用的是均方误差，输出不会被压到 0 或 1，而是接近目标值的小数。判断模型是否学会 XOR，看的是「接近 0 的输入预测值小、接近 1 的输入预测值大」这个趋势是否成立。

一个完整示例

把上面的步骤串起来，下面是一份从头到尾可运行的完整脚本。它包含设备选择、模型定义、数据准备、训练与评估：

import torch
import torch.nn as nn
import torch.optim as optim

# 1. 选择设备
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f'Using device: {device}')

# 2. 定义模型
class SimpleNN(nn.Module):
    def __init__(self):
        super(SimpleNN, self).__init__()
        self.fc1 = nn.Linear(2, 4)
        self.fc2 = nn.Linear(4, 1)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# 3. 准备数据
X_train = torch.tensor([[0.0, 0.0], [0.0, 1.0], [1.0, 0.0], [1.0, 1.0]], device=device)
y_train = torch.tensor([[0.0], [1.0], [1.0], [0.0]], device=device)

# 4. 实例化模型、损失函数、优化器
model = SimpleNN().to(device)
criterion = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.1)

# 5. 训练
for epoch in range(100):
    model.train()
    outputs = model(X_train)
    loss = criterion(outputs, y_train)

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    if (epoch + 1) % 10 == 0:
        print(f'Epoch [{epoch + 1}/100], Loss: {loss.item():.4f}')

# 6. 评估
model.eval()
with torch.no_grad():
    predictions = model(X_train)
    print(f'Predictions: \n{predictions}')

与前面分步代码相比，这份脚本多了两处 device 处理：数据张量在创建时指定 device，模型用 .to(device) 搬到同一设备。PyTorch 要求参与同一次运算的张量位于同一设备上，模型在 GPU 而数据在 CPU 会直接报错，这是初学者最常见的错误之一。

用 Dataset 与 DataLoader 组织数据

上面的例子把全部数据一次性喂给模型。真实任务的数据量往往很大，需要分批、打乱地迭代，这正是 Dataset 与 DataLoader 的职责。

自定义 Dataset

继承 torch.utils.data.Dataset，实现三个方法：__init__ 负责加载数据，__len__ 返回样本总数，__getitem__ 按索引返回单个样本：

import torch
from torch.utils.data import Dataset, DataLoader

class MyDataset(Dataset):
    def __init__(self):
        self.data = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
        self.labels = torch.tensor([0, 1, 0])

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        return self.data[idx], self.labels[idx]

dataset = MyDataset()
dataloader = DataLoader(dataset, batch_size=2, shuffle=True)

for batch in dataloader:
    print("Batch Data:", batch[0])
    print("Batch Labels:", batch[1])

batch_size=2 表示每批取 2 个样本，shuffle=True 表示每个 epoch 开始前打乱顺序。打乱能避免模型学到样本排列带来的虚假规律。由于样本总数是 3、批大小是 2，最后一个批次只会有 1 个样本，这是正常现象。

用 DataLoader 驱动训练

把 DataLoader 接进训练循环，只需把「一次算全部数据」换成「逐批计算」：

for epoch in range(2):
    for inputs, labels in dataloader:
        outputs = inputs + 1
        print(f"Epoch {epoch + 1}, Inputs: {inputs}, Labels: {labels}, Outputs: {outputs}")

这段代码用 inputs + 1 代替真实的前向传播，目的是展示批次迭代的结构。实际训练时，把 outputs = inputs + 1 换成 outputs = model(inputs)，再补上损失计算、梯度清零、反向传播与参数更新四步即可。分批处理的好处是计算效率更高，在 GPU 等硬件加速器上尤其明显。

图像数据增强

如果任务是图像分类，torchvision.transforms 提供了常用的随机变换，比如随机水平翻转、转成张量。这些变换能增加数据的多样性，帮助模型在未见过的数据上泛化得更好：

import torchvision.transforms as transforms
from PIL import Image

image = Image.open('example.jpg')  # 替换为你自己的图片文件

transform = transforms.Compose([
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor()
])

augmented_image = transform(image)
print("Augmented Image Shape:", augmented_image.shape)

transforms.Compose 把多个变换按顺序串起来：先随机水平翻转，再转成张量。ToTensor() 会把像素值缩放到 0 到 1 之间，并把通道维度放到最前面，因此打印出的形状是「通道 × 高 × 宽」。运行前需要准备一张真实存在的图片，并把文件名替换掉。

进阶模型方向

掌握上面的流程后，可以按任务类型换用不同的网络层，训练循环的骨架基本不变。

卷积神经网络（CNN）

处理图像时用 torch.nn.Conv2d 做卷积、配合池化层降采样。在卷积层之后接 torch.nn.BatchNorm2d 做批归一化，可以稳定训练过程、加快收敛。批归一化对卷积层的输出做归一化，是图像模型里的常见配置。

循环神经网络（RNN）

处理序列数据时用 torch.nn.LSTM 或 torch.nn.GRU。这两类模块在 PyTorch 里可以直接调用，适合文本、时间序列等有先后顺序的任务。它们同样继承自 nn.Module，可以像 nn.Linear 一样嵌进自定义模型里。

注意事项

• 梯度必须清零：PyTorch 默认累加梯度，每轮训练前调用 optimizer.zero_grad()，否则参数更新会基于错误的梯度。
• 设备必须一致：模型与输入张量要在同一设备上。用 .to(device) 统一搬运，不要一部分在 CPU、一部分在 GPU。
• 训练与评估模式要切换：训练时 model.train()，评估时 model.eval()，并在评估时用 torch.no_grad() 关闭梯度。
• 形状要对齐：view 要求元素总数不变；模型输出与标签的形状要匹配，否则损失计算会报错或广播出错误结果。
• 显存有限：GPU 显存是硬约束，张量过大或批大小过大会触发显存不足。用 torch.cuda.memory_allocated() 与 torch.cuda.memory_reserved() 观察占用情况。
• 版本信息会变：CUDA 版本、安装命令、各版本支持情况随发布更新，安装前请以 PyTorch 官网当前信息为准，不要直接照搬本文示例中的版本号。
• 损失值不是准确率：均方误差输出的是连续值，判断分类效果需要自己设定阈值或换用合适的损失函数，不能只看损失数字。