目录
训练(training)阶段
会包含三步:
- 前向传播(forward):算预测结果
- 计算损失(loss)
- 反向传播(backward):计算梯度
- 参数更新(optimizer.step)
👉 这一步才需要梯度、计算图、反传。
在 PyTorch 里通常写成:
python
import torch
import torch.nn as nn
import torch.optim as optim
# 1. 定义模型
model = nn.Linear(10, 1)
# 2. 损失函数 & 优化器
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3)
# 3. 切换到训练模式
model.train()
for epoch in range(100):
# 假数据
x = torch.randn(32, 10)
y = torch.randn(32, 1)
# ===== 前向传播 =====
y_pred = model(x)
loss = criterion(y_pred, y)
# ===== 反向传播 =====
optimizer.zero_grad() # 清空旧梯度
loss.backward() # 计算梯度
optimizer.step() # 更新参数
print(f"epoch {epoch}, loss = {loss.item():.4f}")
推理(inference)阶段
只做:
- ✅ 前向传播
- ❌ 不计算 loss
- ❌ 不做反向传播
- ❌ 不更新参数
因此:
- 不需要梯度
- 不需要保留计算图
- 显存占用和计算量都更小
- 速度更快
在 PyTorch 里通常写成:
python
import torch
import torch.nn as nn
# 已训练好的模型
model = nn.Linear(10, 1)
model.load_state_dict(torch.load("model.pt"))
# 1. 切换到推理模式
model.eval()
# 2. 禁用梯度
with torch.no_grad():
x = torch.randn(1, 10)
# ===== 仅前向传播 =====
y_pred = model(x)
print(y_pred)