使用Jupyter lab测试执行代码验证。
PyTorch
PyTorch 是最流行的深度学习框架之一,由 Meta(Facebook)开发。它用张量(Tensor)处理数据,支持 GPU 加速,提供自动求导和构建神经网络的工具。
安装:
sh
pip install torch torchvision
torchvision 是 PyTorch 的视觉工具包,它提供常用数据集和图像预处理。
Tensor
Tensor(张量)是 PyTorch 的核心数据结构,和 NumPy 的 ndarray 类似,但多了两个能力:
js
NumPy ndarray → CPU 上运算
PyTorch Tensor → CPU 或 GPU 上运算,支持自动求导
Tensor ≈ 带 GPU 加速和微分能力的 NumPy 数组。
创建 Tensor
py
import torch
# 从列表创建
t = torch.tensor([1, 2, 3, 4, 5])
print(t)
print(t.dtype)
# 指定类型
t = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32)
print(t.dtype)
# 创建全 0 / 全 1
t = torch.zeros(3, 3)
t = torch.ones(3, 3)
# 随机数
t = torch.rand(3, 3) # 0-1 均匀分布
t = torch.randn(3, 3) # 标准正态分布
# 创建范围
t = torch.arange(0, 10, 2)
Tensor 运算
py
a = torch.tensor([1.0, 2.0, 3.0])
b = torch.tensor([4.0, 5.0, 6.0])
# 基本运算
print(a + b)
print(a * b)
print(a @ b)
# 形状操作
t = torch.rand(2, 3)
print(t.shape)
# 变成任意形状
print(t.reshape(3, 2))
# 转置,行列互换,只能用于二维
print(t.T)
@ - 矩阵乘法, 行x列求和。如果是一维数组,向量对应相乘再求和返回一个标量。 * - 逐元素乘法。
Tensor 与 NumPy 互转
py
import numpy as np
# NumPy → Tensor
arr = np.array([1, 2, 3])
t = torch.from_numpy(arr)
# Tensor → NumPy
arr = t.numpy()
GPU 加速
py
# 自动选择可用设备
if torch.cuda.is_available():
# NVIDIA GPU / Linux/ Windows
device = "cuda"
elif torch.backends.mps.is_available():
# Apple GPU M1 / M2 / M2 / M3
device = "mps"
else:
device = "cpu"
print(f"Using device: {device}")
t = torch.tensor([1, 2, 3])
t_gpu = t.to(device)
t_cpu = t_gpu.to("cpu")
训练模型数据时启用 GPU 加速,模型和数据都放GPU;结果操作时启用 CPU。
自动求导 Autograd
深度学习的核心是反向传播 ------计算损失函数对每个参数的梯度,然后更新参数。Autograd 自动完成这个过程。
梯度的作用让模型预测的结果更准确。趋近于0的梯度,当前参数就是最优解。
模型参数需要requires_grad=True,输入数据不需要。
py
# 创建一个 tensor,开启梯度追踪
x = torch.tensor(2.0, requires_grad=True)
# 前向计算
y = x ** 2 + 3 * x + 1 # y = x² + 3x + 1
# 反向传播:计算 dy/dx 变化率
y.backward()
# 梯度 = dy/dx = 2x + 3 = 2*2 + 3 = 7
print(x.grad)
循环调整不同的参数来计算梯度,找到最优解。
py
# 1. 固定轮数,就跑十次
x = torch.tensor(2.0, requires_grad=True)
for i in range(10):
y = x ** 2 + 3 * x + 1
y.backward()
with torch.no_grad():
x -= 0.1 * x.grad
x.grad.zero_()
# 2. 相对变化率
x = torch.tensor(2.0, requires_grad=True)
prev_loss = float("inf")
while True:
y = x ** 2 + 3 * x + 1
loss = y.item()
if abs(prev_loss - loss)/prev_loss < 0.001:
print(f"损失不再下降,停止训练")
break
prev_loss = loss
y.backward()
with torch.no_grad():
x -= 0.1 * x.grad
x.grad.zero_()
# 3. 早停
x = torch.tensor(2.0, requires_grad=True)
best_loss = float("inf")
patience = 10 # 连续10 轮没改善就停
counter = 0
for i in range(1000):
y = x ** 2 + 3 * x + 1
loss = y.item()
if loss < best_loss - 0.001: # 改善超过0.001,更新参数
best_loss = loss
counter = 0
else:
counter += 1
if counter >= patience:
print(f"连续{patience} 轮没改善,停止训练")
break
y.backward()
with torch.no_grad():
x -= 0.1 * x.grad
x.grad.zero_()
y = x ** 2 + 3 * x + 1
y.backward()
print(x.grad)
在调整x的时候,有一个参数0.1 是学习率,用来控制模型参数的更新。
为什么要一直调用x.grad.zero_(), 因为训练过程中,每次计算梯度,都会将梯度累加到x.grad中,所以每次训练完一个epoch,需要将梯度清零。
使用torch.optim包中的优化器,如torch.optim.SGD,torch.optim.Adam。不用手动去写更新逻辑。
SGD - 随机梯度下降,推荐学习率 0.01 Adam - 自适应学习率优化器,推荐学习率 0.001
py
x = torch.tensor(2.0, requires_grad=True)
# SGD - 随机梯度下降
optimizer = torch.optim.SGD([x], lr=0.1)
# Adam - 自适应学习率优化器
# optimizer = torch.optim.Adam([x], lr=0.001)
best_loss = float("inf")
patience = 10 # 连续10 轮没改善就停
counter = 0
for i in range(1000):
y = x ** 2 + 3 * x + 1
loss = y.item()
if loss < best_loss - 0.001: # 改善超过0.001,更新参数
best_loss = loss
counter = 0
else:
counter += 1
if counter >= patience:
print(f"连续{patience} 轮没改善,停止训练")
break
optimizer.zero_grad()
y.backward()
optimizer.step()
为了找到最优的学习率lr, 可以使用学习率调度器torch.optim.lr_scheduler。 或者通过枚举寻找最优的lr。
学习率的影响:
text
学习率太大:损失震荡或爆炸,训练不稳定
学习率太小:损失下降太慢,训练时间长
学习率合适:损失平稳下降,收敛到最优
构建神经网络
nn.Module
nn.Module 是 PyTorch 中构建神经网络的基类,所有模型都继承它。
nn.Module 提供了构建神经网络的基本功能,如定义层,定义数据流,定义参数。
定义层- 创建层,如全连接层、卷积层、池化层。用来处理输入数据。定义数据流- 定义数据流,如前向传播,反向传播。用来定义如何处理输入数据。定义参数- 定义参数,如权重、偏置。用来保存模型参数。
py
import torch.nn as nn
class SimpleNet(nn.Module):
def __init__(self):
super().__init__()
# 定义层
self.fc1 = nn.Linear(10, 64) # 输入 10,输出 64
self.fc2 = nn.Linear(64, 32) # 输入 64,输出 32
self.fc3 = nn.Linear(32, 1) # 输入 32,输出 1
def forward(self, x):
# 定义数据流
x = self.fc1(x)
x = self.fc2(x)
x = self.fc3(x)
return x
# 创建模型
model = SimpleNet()
# 使用模型
x = torch.randn(10)
y = model(x)
# 输出模型结构
print(model)
print(y)
# 查看模型参数
for name, param in model.named_parameters():
print(name, param.shape)
nn.Module 的子类必须实现 __init__ 和 forward 方法。__init__ 方法用于定义层,forward 方法用于定义数据流。
调用模型实例时,会自动调用forward 方法。数据流会从输入层开始,经过每一层(每一层都进行矩阵乘法),最后输出结果。
激活函数
没有激活函数,多层网络和单层等价。激活函数引入非线性,让网络能学习复杂规律。
上一节的三层网络没有激活函数,等价于单层,没有发挥多层的优势。
ReLU- 最常用,负数变 0,正数不变。隐藏层首选。Sigmoid-适合二分类。任意数字压缩到0-1之间,0输出0.5;越大越接近1;越小越接近0。Softmax- 适合多分类。把任意数字转概率分布,每个数的 e 次方 ÷ 所有数的 e 次方之和Tanh- 适合隐藏层。把任意数字压缩到-1-1之间,0输出0;越大越接近1;越小越接近-1。
py
# ReLU
relu = nn.ReLU()
print(relu(torch.tensor([-1.0, 0.0, 1.0])))
# Sigmoid
sigmoid = nn.Sigmoid()
print(sigmoid(torch.tensor([0.0])))
# Softmax
softmax = nn.Softmax(dim=0)
print(softmax(torch.tensor([1.0, 2.0, 3.0])))
# Tanh
tanh = nn.Tanh()
print(tanh(torch.tensor([0.0])))
上一节增加激活函数:
py
class SimpleNet(nn.Module):
def __init__(self):
super().__init__()
# 定义层
self.fc1 = nn.Linear(10, 64) # 输入 10,输出 64
self.fc2 = nn.Linear(64, 32) # 输入 64,输出 32
self.fc3 = nn.Linear(32, 1) # 输入 32,输出 1
self.relu = nn.ReLU()
self.sigmoid = nn.Sigmoid()
def forward(self, x):
# 定义数据流
x = self.relu(self.fc1(x))
x = self.relu(self.fc2(x))
x = self.sigmoid(self.fc3(x))
return x
# 创建模型
model = SimpleNet()
x = torch.randn(10)
print(model(x))
隐藏层一般使用ReLU,输出层看任务,回归不需要加激活;二分类则使用Sigmoid;多分类则使用Softmax。
损失函数
衡量模型预测和真实值的差距。损失越小,模型预测越准确。
- 回归任务使用
MSELoss均方误差。 - 二分类任务使用
BCELoss / BCEWithLogitsLoss二元交叉熵。 - 多分类任务使用
CrossEntropyLoss交叉熵。
根据损失函数告诉模型错误的有多离谱,让模型调整参数。
py
# 回归
mse = nn.MSELoss()
loss = mse(torch.tensor([1.0, 2.0]), torch.tensor([1.5, 2.5]))
print(loss)
# 二分类
bce = nn.BCELoss()
loss = bce(torch.tensor([0.9, 0.1]), torch.tensor([1., 0.]))
print(loss)
# 多分类
ce = nn.CrossEntropyLoss()
loss = ce(torch.tensor([[0.1, 0.9], [0.8, 0.2]]),
torch.tensor([1, 0]))
print(loss)
BCELoss 输入必须是0~1的概率,必须是浮点数;输出层要配合Sigmoid。
CrossEntropyLoss 输入是原始分数,内部自动做了Softmax。
在之前的示例中增加损失函数:
py
# ... 省略之前的模型示例代码
# 损失函数
criterion = nn.BCELoss() # 二分类用 BCELoss
# 准备数据
X = torch.randn(5, 10) # 5 个样本,10 个特征
y = torch.tensor([1., 0., 1., 0., 1.]) # 真实标签
# 前向计算
output = model(X) # 输出 5 个概率值
print(f"预测:{output.squeeze()}")
print(f"真实:{y}")
# 计算损失
loss = criterion(output.squeeze(), y)
print(f"损失:{loss.item():.4f}")
DataLoader
DataLoader 是一个迭代器,用于加载数据。它会自动将数据分批次处理,并返回一个迭代器。
py
from torch.utils.data import DataLoader, TensorDataset
# 创建数据集
X = torch.randn(100, 10)
y = torch.randn(100, 1)
dataset = TensorDataset(X, y)
# 创建 DataLoader: batch_size - 批次大小, shuffle - 是否打乱数据
dataloader = DataLoader(dataset, batch_size=16, shuffle=True)
# 遍历
for batch_X, batch_y in dataloader:
print(batch_X.shape)
print(batch_y.shape)
完整的训练流程加载数据 → 前向计算 → 算损失 → 反向传播 → 更新参数 → 重复.
分批次处理可以提升训练效率,减轻模型内存占用。
按照完整的流程:
- 1.创建数据集
py
from torch.utils.data import DataLoader, TensorDataset
# 200 个样本,5 个特征
X = torch.randn(200, 5)
# 定义标签,根据每个样本的前两个值确定标签
y = (X[:, 0] + X[:, 1] > 0).float().unsqueeze(1)
# 创建数据集
dataset = TensorDataset(X, y)
# 创建数据加载器:批次大小为 32,随机打乱
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)
- 2.创建模型
py
# 创建模型: 按顺序执行
model = nn.Sequential(
nn.Linear(5, 16),
nn.ReLU(),
nn.Linear(16, 1),
nn.Sigmoid()
)
- 3.创建损失函数和优化器
py
# 根据 y 标签判断任务是一个二分类任务
criterion = nn.BCELoss()
# 优化器
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
- 4.训练模型
py
best_loss = float('inf')
patience = 10
counter = 0
for epoch in range(1000):
total_loss = 0
for batch_X,batch_y in dataloader:
y_pred = model(batch_X)
loss = criterion(y_pred, batch_y)
# 反向传播
optimizer.zero_grad()
loss.backward()
optimizer.step()
total_loss += loss.item()
avg_loss = total_loss / len(dataloader)
if avg_loss < best_loss:
best_loss = avg_loss
counter = 0
else:
counter += 1
if counter >= patience:
break
- 5.模型评估
py
from sklearn.metrics import classification_report
test_X = torch.randn(20, 5)
test_y = (test_X[:, 0] + test_X[:, 1] > 0).float().unsqueeze(1)
model.eval()
with torch.no_grad():
test_pred = model(test_X)
# 二分类
predicted = (test_pred > 0.5).float()
accuracy = (predicted == test_y).float().mean()
print(f'Accuracy: {accuracy.item():.4f}')
# 混淆矩阵
print(classification_report(test_y.numpy(), predicted.numpy()))
- 6.模型保存
py
# 保存模型
torch.save(model.state_dict(), 'model.pt')
# 加载模型:和训练模型结构必须保持一致
model = nn.Sequential(
nn.Linear(5, 16),
nn.ReLU(),
nn.Linear(16, 1),
nn.Sigmoid()
)
model.load_state_dict(torch.load('model.pt'))
步骤5、6是使用层面的,这里也演示下后续训练完模型后,模型评估、模型保存和加载。
CNN 卷积神经网络
用途: 图像识别、图像分类
原理: 用卷积核扫描图像,提取局部特征(边缘、纹理、形状)。
输入图像 → 卷积层(提取特征)→ 池化层(压缩)→ 全连接层(分类)
- 1.定义输入图像数据
py
# 32x32 灰度图像
# 1张,1通道,32x32
x = torch.randn(1, 1, 32, 32)
- 2.定义卷积层 - 提取特征
py
# 1 输入通道,16输出通道,卷积核大小3x3
conv = nn.Conv2d(in_channels=1, out_channels=16, kernel_size=3)
# 处理数据
output = conv(x)
print(output.shape)
输入通道由数据决定;输出通过(卷积核数量)自定义,越大特征越多,处理也就越慢;卷积核定义特征检测器,在图像上滑动提取局部特征。
16个卷积核对应不同的特征权重,提取不同的特征。3x3的提取器处理32x32的图像,输出30x30,滑动提取只能到[29,29]导致图片的大小丢失2x2,输出为30x30。
怎么解决图像大小的问题?可以通过配置补充padding来让提取器多滑动两次。最后输出原图大小32x32
py
conv = nn.Conv2d(in_channels=1, out_channels=16, kernel_size=3,padding=1)
- 3.定义池化层 - 压缩
压缩特征图,只保留重要信息。特定大小的区域只取最显著的特征或者是平均值。
py
# 最大池化:取 2x2 区域的最大值
pool = nn.MaxPool2d(kernel_size=2)
result = pool(torch.relu(output))
print(result.shape)
# 平均池化:取 2x2 区域的平均值
pool = nn.AvgPool2d(kernel_size=2)
result = pool(torch.relu(output))
print(result.shape)
- 4.定义全连接层 - 分类
把特征变成分数。
py
# 展平
output = output.view(output.size(0), -1)
print(output.shape)
# 全连接层
linear = nn.Linear(in_features=output.size(1), out_features=128)
result = torch.relu(linear(output))
linear = nn.Linear(128, out_features=10)
result = linear(result)
print(result.shape)
完整的CNN卷积模型:
py
class CNN(nn.Module):
def __init__(self):
super(CNN, self).__init__()
self.conv = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3,padding=1)
self.pool = nn.MaxPool2d(2, 2)
self.fc1 = nn.Linear(32 * 16 * 16, 128)
self.fc2 = nn.Linear(128, 10)
def forward(self, x):
x = self.pool(torch.relu(self.conv(x)))
x = x.view(-1, 32 * 16 * 16)
x = torch.relu(self.fc1(x))
x = self.fc2(x)
return x
# 测试
model = CNN()
x = torch.randn(10,1,32,32)
print(model(x).shape)
RNN / LSTM
用途: 序列数据(文本、时间序列、语音)
原理: RNN 有"记忆",能记住前面的信息影响后面的输出。
RNN:是短期记忆;LSTM:是长期记忆,有选择地记住和遗忘
RNN 适合序列少、数据量少、速度快的;LSTM 适合长序列、数据充足、任务复杂。
RNN
定义RNN:
input_size - 输入向量的维度 hidden_size - 隐藏层向量的维度 num_layers - 堆叠的 RNN 层数 batch_first - 输入和输出 Tensor 的第一个维度是批量维度(batch)
py
# 定义 RNN
rnn = nn.RNN(input_size=10, hidden_size=20, num_layers=1, batch_first=True)
# 输入 batch_size, seq_len, input_size
x = torch.randn(32, 5, 10)
# 前向计算
output, h_n = rnn(x)
print(output.shape)
print(h_n.shape)
结果输出output是每个时间步的输出;h_n是最终隐藏状态(包含整个序列信息)
LSTM
LSTM 和RNN 类似,只是 LSTM 多了一个细胞状态 c_n。
py
# 定义 LSTM
lstm = nn.LSTM(input_size=10, hidden_size=20, num_layers=1, batch_first=True)
x = torch.randn(32, 5, 10)
# 前向计算
output, (h_n, c_n) = lstm(x)
print(output.shape)
print(h_n.shape)
print(c_n.shape)
output是每个时间步的输出;h_n是最终隐藏状态(包含整个序列信息); c_n是最终细胞状态,通过它实现长期记忆;
Transformer
现代大语言模型(GPT、BERT)的基础架构。RNN/LSTM是逐步处理,一个接一个慢 ;Transformer 是并行 处理,多个时间步同时处理,快
核心思想:Self-Attention 自注意力
让每个位置都能"看到"序列中所有其他位置,捕捉全局关系。
js
输入序列 → Self-Attention(每个词关注其他词)→ 前馈网络 → 输出
↑
重复 N 次
Self-Attention 计算过程
每个词生成三个向量:
js
Q(Query):我要找什么信息
K(Key): 我有什么特征
V(Value):我的具体内容
注意力 = softmax(Q × K^T / √d) × V
RNN 逐步处理,天然有位置信息;Transformer 并行处理,不知道词的顺序,需要给每个词加上位置编码。
在PyTorch 中使用 Transformer,提供了两种模式:Encoder编码器 和 Decoder解码器。
Encoder 可以双向理解输入序列,并生成输出序列。理解更准确,BERT 模型就是 Encoder。
py
import torch.nn as nn
# Transformer 编码器层
encoder_layer = nn.TransformerEncoderLayer(
d_model=512, # 特征维度
nhead=8, # 注意力头数
dim_feedforward=2048, # 前馈网络隐藏层维度
dropout=0.1, # 丢弃概率
batch_first=True,
)
# 堆叠 N 层
transformer = nn.TransformerEncoder(encoder_layer, num_layers=6)
# 输入:(seq_len, batch, d_model)
src = torch.rand(10, 32, 512)
output = transformer(src)
print(output.shape)
Decoder 是生成输出。只能从左往右输入理解。
py
# Decoder 层
decoder_layer = nn.TransformerDecoderLayer(
d_model=512,
nhead=8,
dim_feedforward=2048,
dropout=0.1,
batch_first=True
)
# 堆叠 6 层
decoder = nn.TransformerDecoder(decoder_layer, num_layers=6)
# 输入
memory = torch.rand(32, 10, 512) # Encoder 的输出(理解结果)
tgt = torch.rand(32, 8, 512) # 目标序列(要生成的内容)
# 训练
output = decoder(tgt, memory)
print(output.shape)
# 使用(推理)
decoder.eval()
with torch.no_grad():
# 假设 tgt 是一个 batch 的输入序列
output = decoder(tgt, memory)
print(output.shape)
memory 是 Encoder 的输出(理解输入),tgt 是目标序列(要生成的内容)。Decoder 根据 memory 和 tgt 生成最终结果。
Decoder-Only 架构,现代大模型(如 GPT 系列)主要采用 Decoder-Only 架构。它通过庞大的参数量和海量数据训练,仅用 Decoder 部分就能实现强大的语言理解和生成能力,无需额外的 Encoder。
Hugging Face
Hugging Face 是 AI 领域最重要的开源社区,提供了海量预训练模型和工具。相当于前端的 npm + GitHub + Vercel。
text
Hub → 模型仓库(类比 npm registry)
Datasets → 数据集平台
Spaces → Demo 部署(类比 Vercel)
Transformers → Python SDK(类比 npm install + import)
前面学习了从零搭建和训练 CNN、RNN、Transformer。但在实际项目中,很少从零训练大模型------
| 维度 | 从零训练 | 使用预训练模型 |
|---|---|---|
| 数据量 | 百万级以上 | 几百到几万条即可 |
| 算力 | 数十块 GPU,数周 | 单张消费级显卡 |
| 成本 | 数万到数百万美元 | 几乎为零 |
| 效果 | 依赖数据质量和调参 | 站在巨人的肩膀上 |
回忆上一节讲的 Decoder-Only Transformer(GPT 架构),它就是所有现代大语言模型的基座。Hugging Face 上托管了数万个这样的模型,直接下载就能用。
Hub:下载/上传模型
CLI 工具
在MAC上安装:
sh
brew install hf
# or
# curl -LsSf https://hf.co/cli/install.sh | bash
通过hf --help查看所有命令。 可以通过个人账号创建私有仓库,保存私有模型、文件等。
Python API
sh
pip install huggingface-hub
在程序中使用,下载一个模型:
py
from huggingface_hub import snapshot_download
# 下载整个模型文件夹
snapshot_download("google/bert_uncased_L-2_H-128_A-2", local_dir="./models/bert")
下载后的模型在 ./models/bert/ 目录下,包含 config.json、model.safetensors(权重)、tokenizer.json(分词器)等文件。
Datasets
安装:
sh
pip install datasets
datasets 模块提供了丰富的数据集,可以快速加载。
py
from datasets import load_dataset
# 加载 Hugging Face 上的公开数据集
dataset = load_dataset("stanfordnlp/imdb") # IMDB 电影评论
print(dataset)
print(dataset["train"][0])
# 加载本地数据
dataset = load_dataset("csv", data_files="train.csv")
print(dataset)
pipeline()
安装:
sh
pip install transformers
pipeline() 是 HF 提供的最高层 API,把模型加载、分词、推理全部封装好。几行代码就能跑通一个完整任务,适合先体验再深入。
py
from transformers import pipeline
# 情感分析
classifier = pipeline("sentiment-analysis")
print(classifier("I hate this!"))
# 文本生成
generator = pipeline("text-generation", model="gpt2")
print(generator("I previously", max_length=50)[0]["generated_text"])
pipeline() 背后做了什么?
text
pipeline("sentiment-analysis")
→ 自动选默认模型(distilbert-base-uncased-finetuned-sst-2-english)
→ 下载模型 + 分词器
→ 输入文本 → tokenize → 模型推理 → 输出标签
一行代码背后是完整的加载 → 预处理 → 推理 → 后处理 流水线。理解了这些,就可以用下面的 AutoModel API 获得完全控制权。
| 任务 | pipeline 名称 | 输入 → 输出 |
|---|---|---|
| 情感分析 | sentiment-analysis |
文本 → 正/负面 |
| 文本生成 | text-generation |
提示词 → 续写文本 |
| 文本分类 | text-classification |
文本 → 类别标签 |
| 命名实体识别 | ner |
文本 → 人名/地名/组织 |
| 完形填空 | fill-mask |
带 [MASK] 的句子 → 补全 |
| 图像分类 | image-classification |
图片 → 类别 |
默认的模型可能仅支持英文输入,如果输入中文,需要指定支持中文的模型
实战:MNIST 手写数字识别
把前面学的知识串起来,用 CNN 识别手写数字(0-9)。
MNIST 是深度学习界的 "Hello World":60,000 张 28×28 的灰度手写数字图片,0 到 9 各写法的笔迹,训练模型认出每个数字。
1. 加载数据
下载数据:
py
from torchvision import datasets,transforms
# 预处理:转成 Tensor + 标准化
transform = transforms.Compose([
transforms.ToTensor(), # PIL 图片 → Tensor,同时把 0-255 缩到 0-1
transforms.Normalize((0.1307,), (0.3081,)) # 标准化,让数据分布更稳定
])
# 下载 MNIST(第一次会下载,之后用缓存)
train_dataset = datasets.MNIST("./data", train=True, download=True, transform=transform)
test_dataset = datasets.MNIST("./data", train=False, transform=transform)
print(f"训练集大小:{len(train_dataset)} 张")
print(f"测试集大小:{len(test_dataset)} 张")
print(f"一张图的 shape:{train_dataset[0][0].shape}")
print(f"对应标签:{train_dataset[0][1]}")
datasets 提供的数据是.gz压缩的二进制文件,存储的是原始像素数组。在下载过程中,transforms 提供了对数据进行预处理的方法,进行格式转换、归一化操作。
0.1307 和0.3081 是前人算好的均值和标准差,直接用。
进行数据分割,60,000 张数据很多,分批次训练。
py
from torch.utils.data import DataLoader
# DataLoader 把数据切成小批(batch),每次喂 64 张图给模型
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=1000)
2. 定义 CNN
前面的 CNN 章节:卷积提取特征 → 池化压缩尺寸 → 全连接分类。
py
import torch.nn as nn
class CNN(nn.Module):
def __init__(self):
super().__init__()
# 卷积层:提取图像特征
self.conv1 = nn.Conv2d(1, 32, 3, 1, padding=1) # 输入 1 通道(灰度),输出 32 通道
self.conv2 = nn.Conv2d(32, 64, 3, 1, padding=1) # 32 → 64 通道
# 全连接层:分类
self.fc1 = nn.Linear(12544, 128) # 展平后 64×14×14 = 12544
self.fc2 = nn.Linear(128, 10) # 输出 10 类(数字 0-9)
def forward(self, x):
x = torch.relu(self.conv1(x)) # conv1 + ReLU
x = torch.relu(self.conv2(x)) # conv2 + ReLU
x = torch.max_pool2d(x, 2) # 2×2 池化,
x = torch.flatten(x, 1) # 展平成一维向量
x = torch.relu(self.fc1(x)) # 全连接 + ReLU
x = self.fc2(x) # 输出 10 个分数(不加激活,Loss 内置 Softmax)
return x
model = CNN()
print(model)
数据在模型里怎么流的?
以 batch_size=64 为例,DataLoader 一次喂 64 张图,第一维就是 batch:
text
输入 [64, 1, 28, 28] 64 张灰度图,每张 1 通道 28×28
conv1 [64, 32, 28, 28] 每张图 32 个特征图(padding补充,尺寸不变)
conv2 [64, 64, 28, 28] 每张图 64 个特征图
max_pool [64, 64, 14, 14] 2×2 池化
flatten [64, 12544] 每张图 64×14×14=12544 拍平
fc1 [64, 128] 每张图一个中间向量
fc2 [64, 10] 每张图 10 个分数,argmax 取最大的就是预测数字
3. 训练
py
import torch.optim as optim
criterion = nn.CrossEntropyLoss() # 多分类损失函数
optimizer = optim.Adam(model.parameters(), lr=0.001) # Adam 优化器
for epoch in range(5):
model.train() # 切换到训练模式
total_loss = 0
for batch_X, batch_y in train_loader:
optimizer.zero_grad() # 清空上一轮梯度
output = model(batch_X) # 前向传播
loss = criterion(output, batch_y) # 算损失
loss.backward() # 反向传播,算梯度
optimizer.step() # 更新参数
total_loss += loss.item() * batch_X.size(0) # 乘以这个 batch 的样本数
print(f"Epoch {epoch+1}, Avg Loss: {total_loss / len(train_dataset):.4f}")
输出:
text
Epoch 1, Avg Loss: 0.1170
Epoch 2, Avg Loss: 0.0354
Epoch 3, Avg Loss: 0.0215
Epoch 4, Avg Loss: 0.0150
Epoch 5, Avg Loss: 0.0124
Loss 每轮都在下降,说明模型在学会分辨数字。每一轮做的五件事就是训练循环那节讲的核心流程:zero_grad → forward → loss → backward → step。
4. 评估
py
model.eval() # 切换到评估模式
correct = 0
total = 0
with torch.no_grad(): # 不计算梯度,省显存、加速
for batch_X, batch_y in test_loader:
output = model(batch_X)
pred = output.argmax(dim=1) # 取分数最高的那个类别
correct += (pred == batch_y).sum().item()
total += batch_y.size(0)
print(f"测试准确率:{correct / total:.4f}")
输出:
text
测试准确率:0.9876
argmax(dim=1) ------ fc2 输出 10 个分数 [0.1, 0.05, 0.8, ...],取最大值的位置就是预测的数字。
5. 保存模型
py
torch.save(model.state_dict(), "mnist_cnn.pth")
# 下次加载
model.load_state_dict(torch.load("mnist_cnn.pth"))
6. 真实场景测试
准备了一张手写的图片,使用模型识别一下。

py
from PIL import Image, ImageOps
# 加载你的图片,转灰度、缩放到 28×28
img = Image.open("test-eight.jpg").convert("L").resize((28, 28))
img = ImageOps.invert(img) # 黑白反转,变黑底白字
# 转成和训练数据一样的格式
img_tensor = transforms.ToTensor()(img) # [1, 28, 28],0-1
img_tensor = transforms.Normalize((0.1307,), (0.3081,))(img_tensor) # 标准化
# 预测
with torch.no_grad():
pred = model(img_tensor.unsqueeze(0)).argmax(dim=1).item()
print(f"识别结果: {pred}")
识别结果: 8 , 看来写的还是太标准了,没有什么难度。
常见坑
| 坑 | 说明 | 解决 |
|---|---|---|
忘记 optimizer.zero_grad() |
梯度累加 | 每次训练前清空 |
忘记 model.train() / model.eval() |
Dropout/BatchNorm 行为不同 | 训练用 train,评估用 eval |
| GPU 内存不足 | 数据太大 | 减小 batch_size |
| 梯度消失/爆炸 | 网络太深 | 用 ReLU、残差连接、梯度裁剪 |
| 过拟合 | 训练好测试差 | Dropout、数据增强、早停 |