Pytorch官网神经网络Demo详细注释/解释

1、网站地址

Quickstart — PyTorch Tutorials 2.2.1+cu121 documentation

2、详细注释/解释

# 1>>>导入各自依赖包
# PyTorch 有两个函数可以处理数据：torch.utils.data.DataLoader 和 torch.utils.data.Dataset。
# Dataset 存储样本及其相应的标签；DataLoader 将Dataset封装成可迭代的对象。
import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensor
# 2>>>下载训练和测试数据
# train set: 60000; test set:10000; 10 classes
training_data = datasets.FashionMNIST(
    root="data",  # 指定数据集的存储路径。
    train=True,  # 当设置为True时，表示加载训练集；设置为False时，表示加载测试集。
    download=True,  # 当设置为True时，如果数据集不存在于指定的root路径下，就会从互联网上下载数据集。
    transform = ToTensor()  # ToTensor()是将PIL Image或者NumPy ndarray转换为torch.FloatTensor，维度：(C x H x W), 标准化：[0.0, 1.0]
)
test_data = datasets.FashionMNIST(
    root="data",
    train=False,
    # download=True,
    transform=ToTensor()
)

# 3>>>数据批量处理
# 通过将数据集传递给DataLoader（支持自动批处理、采样、洗牌和多进程加载的可迭代对象），
# 每次迭代将返回64个数据样本及其对应标签，以便进行模型训练

# 定义批大小
batch_size = 64
# 创建训练数据的 DataLoader
train_loader = DataLoader(
    dataset=training_data,  # 训练数据集变量名
    batch_size=batch_size, # 每次取得样本大小
    shuffle=True,  # 在每个epoch开始时，对数据进行洗牌
    num_workers=2  # 使用两个进程来加载数据，可以根据您的系统配置进行调整
)

# 创建测试数据的 DataLoader
test_loader = DataLoader(
    dataset=test_data,  # 这里应该是之前定义的测试数据集变量名
    batch_size=batch_size, # 每次取得样本大小
    shuffle=False,  # 通常在测试数据上不需要随机
    num_workers=2  # 使用两个进程来加载数据
)

# 3>>>创建模型
# 使用GPU
device = (
    "cuda"
    if torch.cuda.is_available()  # 检查CUDA是否可用， 可用的话就用CUDA， 反之mps
    else "mps"
    if torch.backends.mps.is_available()  # 检查MPS（Metal Performance Shaders）是否可用
    else "cpu"
)
print(f"Using {device} device")

class NeuralNetwork(nn.Module):  # 定义一个继承自nn.Module的神经网络类
    def __init__(self):  # 类的初始化函数
        super().__init__()  # 调用父类nn.Module的初始化函数
        self.flatten = nn.Flatten()  # 展平为1维张量，28*28 ——> 784
        self.linear_relu_stack = nn.Sequential(   # 使用nn.Sequential构建一个序列化的模块容器
            nn.Linear(28 * 28, 512),  # 第一个线性层（有时也称全连接），将784个输入特征映射到512个输出特征。
            nn.ReLU(),  # 非线性激活函数，用于增加模型的非线性，使得网络能够学习更复杂的函数映射
            nn.Linear(512, 512),  # 第二个线性层，将512个输入特征映射到512个输出特征
            nn.ReLU(),  # 第二个激活函数 
            nn.Linear(512, 10)  # 第三个线性层，将512个输入特征映射到10个输出特征。10代表了输出类别的数量
        )
    def forward(self, x):
        x = self.flatten(x)  # 将输入数据展平，实例具有flatten()方法，因为初始化方法中已定义
        logits = self.linear_relu_stack(x)  # 执行网络
        return logits

# 4>>> 实例化网络
model = NeuralNetwork().to(device) # 使用GPU加速
print(model)

# 5>>>模型参数优化
# 使用损失函数和优化器进行训练优化
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)

# 模型在训练过程中产生预测结果，利用误差反向传播更新模型参数
def train(dataloader, model, loss_fn, optimizer):
    size = len(dataloader.dataset)  # 获取训练数据集size
    model.train()  # 调用pytorch训练方法
    for batch, (X, y) in enumerate(dataloader):  # 分别将数据和标签放入GPU
        X, y = X.to(device),  y.to(device)

        # 计算预测误差
        pred = model(X)  # 数据传输模型训练
        loss = loss_fn(pred, y) # 计算误差

        # 反演传播
        loss.backward()  # 进行反向传播，计算梯度
        optimizer.step()  # 根据计算的梯度更新模型的参数
        optimizer.zero_grad()  # 梯度是累积的，因此每个批次后都需要清零。

        if batch % 100 ==0:  # 每处理100个批次，执行以下打印操作，以便监控训练进度
            loss, current = loss.item(), (batch + 1) * len(X)
            print(f"loss: {loss:>7f}  [{current:>5f} / {size:>5d}; {int(100 * (current / size))}%]")

# 检查测试集得预测效果
def test(dataloader, model, loss_fn):
    size = len(dataloader.dataset)  # 读取数据大小 # test: 10000
    num_batches = len(dataloader)  # 训练得批数 # 10000 / 64 = 157
    model.eval()  # # 调用pytorch测试方法
    test_loss, correct = 0, 0  # 初始化误差和正确个数
    with 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()  # 计算累积误差，tensor转换为数值
            correct += (pred.argmax(1) == y).type(torch.float).sum().item()  # 计算正确个数；argmax(1)在每一行（每个样本）中寻找最大值的索引。
    test_loss /= num_batches  # 计算平均损失
    correct /= size  #  # 计算准确率
    print(f"Test Error: \n Accuracy:{(100 * correct):>0.1f}%, Avg loss:{test_loss:>0.8f}\n")

# 6>>>开始训练
# 模型迭代10次
epochs = 10
for t in range(epochs):
    print(f"{t + 1}\n-------------------------------------")
    train(train_dataloader, model, loss_fn, optimizer)
    test(test_dataloader, model, loss_fn)
print("Done!")

# 保存模型
torch.save(model.state_dict(), "Demo_model.pth")
print("Saved PyTorch Model State to model.pth")

3、运行结果

NeuralNetwork(
  (flatten): Flatten(start_dim=1, end_dim=-1)
  (linear_relu_stack): Sequential(
    (0): Linear(in_features=784, out_features=512, bias=True)
    (1): ReLU()
    (2): Linear(in_features=512, out_features=512, bias=True)
    (3): ReLU()
    (4): Linear(in_features=512, out_features=10, bias=True)
  )
)

-------------------------------------
...
...
...
-------------------------------------
10
-------------------------------------
loss: 0.829423  [64.000000 / 60000; 0%]
loss: 0.911203  [6464.000000 / 60000; 10%]
loss: 0.668101  [12864.000000 / 60000; 21%]
loss: 0.877313  [19264.000000 / 60000; 32%]
loss: 0.765406  [25664.000000 / 60000; 42%]
loss: 0.770926  [32064.000000 / 60000; 53%]
loss: 0.848601  [38464.000000 / 60000; 64%]
loss: 0.820338  [44864.000000 / 60000; 74%]
loss: 0.828129  [51264.000000 / 60000; 85%]
loss: 0.783156  [57664.000000 / 60000; 96%]
Test Error: 
 Accuracy:70.9%, Avg loss:0.78746766

Done!

Process finished with exit code 0