PyTorch 学习笔记

环境：python3.8 + PyTorch2.4.1+cpu + PyCharm

参考链接：

快速入门 — PyTorch 教程 2.6.0+cu124 文档

PyTorch 文档 — PyTorch 2.4 文档

快速入门

导入库

import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensor

加载数据集

使用 FashionMNIST 数据集。每个 TorchVision 都包含两个参数：分别是修改样本和标签。

# Download training data from open datasets.
training_data = datasets.FashionMNIST(root="data",     # 数据集存储的位置train=True,      # 加载训练集（True则加载训练集）download=True,   # 如果数据集在指定目录中不存在，则下载（True才会下载）transform=ToTensor(), # 应用于图像的转换列表，例如转换为张量和归一化
)# Download test data from open datasets.
test_data = datasets.FashionMNIST(root="data",train=False,     # 加载测试集（False则加载测试集）download=True,transform=ToTensor(),
)

创建数据加载器

batch_size = 64# Create data loaders.
# DataLoader()：batch_size每个批次的大小，shuffle=True则打乱数据
train_dataloader = DataLoader(training_data, batch_size=batch_size)
test_dataloader = DataLoader(test_data, batch_size=batch_size)for X, y in test_dataloader: # 遍历训练数据加载器，x相当于图片，y相当于标签print(f"Shape of X [N, C, H, W]: {X.shape}")print(f"Shape of y: {y.shape} {y.dtype}")break

创建模型

为了在 PyTorch 中定义神经网络，我们创建一个继承来自 nn.模块。我们定义网络的各层，并在函数中指定数据如何通过网络。要加速作，我们将其移动到 CUDA、MPS、MTIA 或 XPU 等加速器。如果当前加速器可用，我们将使用它。否则，我们使用 CPU。__init__forward

#使用加速器，并打印当前使用的加速器（当前加速器可用则使用当前的，否则使用cpu）
# device = torch.accelerator.current_accelerator().type if torch.accelerator.is_available() else "cpu" # torch2.4.2并没有accelerator这个属性，2.6的才有，所以注释掉
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using {device} device")# 检查 CUDA 是否可用
print("CUDA available:", torch.cuda.is_available())# Define model
class NeuralNetwork(nn.Module):def __init__(self):super().__init__()self.flatten = nn.Flatten()self.linear_relu_stack = nn.Sequential(nn.Linear(28*28, 512),nn.ReLU(),nn.Linear(512, 512),nn.ReLU(),nn.Linear(512, 10))def forward(self, x):x = self.flatten(x)logits = self.linear_relu_stack(x)return logitsmodel = NeuralNetwork().to(device) #torch2.4.2并没有accelerator这个属性，2.6的才有，所以注释掉不用
# model = NeuralNetwork()
print(model)

优化模型参数

要训练模型，我们需要一个损失函数和一个优化器：

loss_fn = nn.CrossEntropyLoss() # 损失函数，nn.CrossEntropyLoss()用于多分类
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3) # 优化器，用于更新模型的参数，以最小化损失函数

'''
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)
优化器用PyTorch 提供的随机梯度下降（Stochastic Gradient Descent, SGD）优化器
model.parameters()：将模型的参数传递给优化器，优化器会根据这些参数计算梯度并更新它们
lr=1e-3：学习率（learning rate），控制每次参数更新的步长
（较大的学习率可能导致训练不稳定，较小的学习率可能导致训练速度变慢）
'''

在单个训练循环中，模型对训练集进行预测（分批提供给它），并且反向传播预测误差以调整模型的参数：

'''
训练模型（单个epoch）
dataloader：数据加载器，用于按批次加载训练数据
model     ：神经网络模型
loss_fn   ：损失函数，用于计算预测值与真实值之间的误差
optimizer ：优化器，用于更新模型参数
'''
def train(dataloader, model, loss_fn, optimizer):size = len(dataloader.dataset)model.train() # 将模型设置为训练模式（启用 dropout 和 batch normalization 的训练行为）for batch, (X, y) in enumerate(dataloader): # 遍历 dataloader 中的每个批次，获取输入 X 和标签 yX, y = X.to(device), y.to(device) # 将数据移动到指定设备（如 GPU 或 CPU）# Compute prediction error# 计算预测损失，同时也是前向传播pred = model(X)         # 模型的预测值，即模型的输出loss = loss_fn(pred, y) # 计算损失：y为实际的类别标签# Backpropagation 反向传播和优化# 梯度清零应在每次反向传播之前执行，以避免梯度累积(先用optimizer.zero_grad())loss.backward()       # 计算梯度optimizer.step()      # 使用优化器更新模型参数optimizer.zero_grad() # 清除之前的梯度（清零梯度，为下一轮计算做准备）# 梯度清零应在每次反向传播之前执行，以避免梯度累积(在计算模型预测值前先用optimizer.zero_grad())if batch % 100 == 0: # 每 100 个批次打印一次损失值和当前处理的样本数量loss, current = loss.item(), (batch + 1) * len(X)print(f"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]")

进度条显示：

如果数据集较大，训练过程可能较慢。可以使用 tqdm 库添加进度条，提升用户体验。例如：
from tqdm import tqdm
for batch, (X, y) in enumerate(tqdm(dataloader, desc="Training")):...

我们还根据测试集检查模型的性能，以确保它正在学习：

# 测试模型
def test(dataloader, model, loss_fn):size = len(dataloader.dataset)  # 测试集的总样本数num_batches = len(dataloader)   # 测试数据加载器（dataloader）的总批次数model.eval()                    # 设置为评估模式，这会关闭 dropout 和 batch normalization 的训练行为test_loss, correct = 0, 0       # 累积测试损失和正确预测的样本数with torch.no_grad(): # 禁用梯度计算，使用 torch.no_grad() 上下文管理器，避免计算梯度，从而节省内存并加速计算for X, y in dataloader:X, y = X.to(device), y.to(device) # 将数据加载到指定设备pred = model(X) # 模型预测test_loss += loss_fn(pred, y).item() # 累积损失correct += (pred.argmax(1) == y).type(torch.float).sum().item() # 累积正确预测数# correct += (pred.argmax(1) == y).float().sum().item()  # 可以直接使用 .float()，更简洁test_loss /= num_batches    # 平均损失correct /= size             # 准确率print(f"Test Error: \n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f} \n")

correct += (pred.argmax(1) == y).type(torch.float).sum().item() # 累积正确预测数
# correct += (pred.argmax(1) == y).float().sum().item()  # 可以直接使用 .float()，更简洁
'''pred.argmax(1):
pred 是模型的输出（通常是未经过 softmax 的 logits，形状为 [batch_size, num_classes]）。
argmax(1) 表示在第二个维度（即类别维度）上找到最大值的索引，返回一个形状为 [batch_size] 的张量，表示每个样本的预测类别。pred.argmax(1) == y:
y 是真实标签（形状为 [batch_size]），表示每个样本的真实类别。
这一步会比较预测的类别和真实类别，返回一个布尔张量，形状为 [batch_size]，其中每个元素表示对应样本的预测是否正确。.type(torch.float):
将布尔张量转换为浮点数张量（True 转为 1.0，False 转为 0.0）.sum():
对浮点数张量求和，得到预测正确的样本总数.item():
将结果从张量转换为 Python 的标量（整数）
'''

举例一：

pred 是模型的输出：torch.tensor([[2.5, 0.3, 0.2], [0.1, 3.2, 0.7]])

y 是真实标签：torch.tensor([0, 1])

import torchpred = torch.tensor([[2.5, 0.3, 0.2], [0.1, 3.2, 0.7]])
y = torch.tensor([0, 1])correct = (pred.argmax(1) == y).type(torch.float).sum().item()
print(correct)  # 输出: 2.0 -> 转换为整数后为 2

举例二：

import torch# 模型输出（未经过 softmax 的 logits）
pred = torch.tensor([[2.0, 1.0, 0.1],  # 第一个样本的预测分数[0.5, 3.0, 0.2],  # 第二个样本的预测分数[1.2, 0.3, 2.5]]) # 第三个样本的预测分数# 真实标签
y = torch.tensor([0, 1, 2])  # 第一个样本的真实类别是 0，第二个是 1，第三个是 2# 计算预测正确的样本数
correct = (pred.argmax(1) == y).type(torch.float).sum().item()
print(f"预测正确的样本数: {correct}") # 预测正确的样本数: 3'''逐步分析
对每个样本的预测分数取最大值的索引，得到预测类别：
pred.argmax(1)  # 输出: tensor([0, 1, 2])比较预测类别和真实标签，得到布尔张量：
pred.argmax(1) == y  # 输出: tensor([True, True, True]).type(torch.float): 将布尔张量转换为浮点数张量：
(pred.argmax(1) == y).type(torch.float)  # 输出: tensor([1.0, 1.0, 1.0]).sum(): 对浮点数张量求和，得到预测正确的样本总数：
(pred.argmax(1) == y).type(torch.float).sum()  # 输出: tensor(3.0).item():将结果从张量转换为 Python 标量：
(pred.argmax(1) == y).type(torch.float).sum().item()  # 输出: 3在这个例子中，模型对所有 3 个样本的预测都正确，因此预测正确的样本数为 3。
'''
# 如果知道总样本数，可以进一步计算准确率：# 总样本数
total = len(y)# 准确率
accuracy = correct / total
print(f"准确率: {accuracy * 100:.2f}%")

训练过程分多次迭代（epoch） 进行。在每个 epoch 中，模型会学习参数进行更好的预测。

然后打印模型在每个 epoch 的准确率和损失，

期望看到准确率Accuracy增加，损失Avg loss随着每个 epoch 的减少而减少：

# 跑5轮，每轮皆是先训练，然后测试
epochs = 5
for t in range(epochs):print(f"Epoch {t+1}\n-------------------------------")train(train_dataloader, model, loss_fn, optimizer)test(test_dataloader, model, loss_fn)
print("Done!")

输出：

Epoch 1
-------------------------------
loss: 2.308106  [   64/60000]
loss: 2.292096  [ 6464/60000]
loss: 2.280747  [12864/60000]
loss: 2.273108  [19264/60000]
loss: 2.256617  [25664/60000]
loss: 2.240094  [32064/60000]
loss: 2.229981  [38464/60000]
loss: 2.204926  [44864/60000]
loss: 2.201917  [51264/60000]
loss: 2.178733  [57664/60000]
Test Error: Accuracy: 46.1%, Avg loss: 2.164820 Epoch 2
-------------------------------
loss: 2.178193  [   64/60000]
loss: 2.160645  [ 6464/60000]
loss: 2.110801  [12864/60000]
loss: 2.129119  [19264/60000]
loss: 2.078400  [25664/60000]
loss: 2.029629  [32064/60000]
loss: 2.044328  [38464/60000]
loss: 1.972220  [44864/60000]
loss: 1.980023  [51264/60000]
loss: 1.920835  [57664/60000]
Test Error: Accuracy: 56.2%, Avg loss: 1.906657 Epoch 3
-------------------------------
loss: 1.938616  [   64/60000]
loss: 1.902610  [ 6464/60000]
loss: 1.797264  [12864/60000]
loss: 1.844325  [19264/60000]
loss: 1.726765  [25664/60000]
loss: 1.688332  [32064/60000]
loss: 1.695883  [38464/60000]
loss: 1.605903  [44864/60000]
loss: 1.628846  [51264/60000]
loss: 1.532240  [57664/60000]
Test Error: Accuracy: 59.8%, Avg loss: 1.541237 Epoch 4
-------------------------------
loss: 1.604458  [   64/60000]
loss: 1.563167  [ 6464/60000]
loss: 1.426733  [12864/60000]
loss: 1.503305  [19264/60000]
loss: 1.376496  [25664/60000]
loss: 1.381424  [32064/60000]
loss: 1.371971  [38464/60000]
loss: 1.312882  [44864/60000]
loss: 1.342990  [51264/60000]
loss: 1.244696  [57664/60000]
Test Error: Accuracy: 62.7%, Avg loss: 1.268371 Epoch 5
-------------------------------
loss: 1.344515  [   64/60000]
loss: 1.318664  [ 6464/60000]
loss: 1.166471  [12864/60000]
loss: 1.275481  [19264/60000]
loss: 1.146058  [25664/60000]
loss: 1.179018  [32064/60000]
loss: 1.171105  [38464/60000]
loss: 1.129168  [44864/60000]
loss: 1.163182  [51264/60000]
loss: 1.077062  [57664/60000]
Test Error: Accuracy: 64.7%, Avg loss: 1.097442 Done!

保存模型

保存模型的常用方法是序列化内部状态字典（包含模型参数）：

torch.save(model.state_dict(), "model.pth")
print("Saved PyTorch Model State to model.pth")

加载模型

加载模型的过程包括重新创建模型结构和加载 state 字典放入其中。

model = NeuralNetwork().to(device)
model.load_state_dict(torch.load("model.pth", weights_only=True))

查看安装的PyTorch版本

方法一：cmd终端查看

终端中输入：

>>>python
>>>import torch
>>>torch.__version__  //注意version前后是两个下划线

方法二：PyCharm查看

打开Pycharm，在Python控制台中输入：

或者在Pycharm的“Python软件包”中查看：