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软件包”中查看：