pytorch 入门

import matplotlib.pyplot as plt
from torchvision.transforms import ToTensor
import torch
from torch import nn  # 包含构建神经网络的所有模块
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

# 训练数据集
training_data = datasets.FashionMNIST(
    root="./data",          # 存储测试集和训练集的路径
    train=True,           # 训练集
    download=True,        # 如果本机没有数据集，就会下载到 root 目录下。
    transform=ToTensor()  # 对样本数据进行处理，转换为张量数据
)
# 测试数据集
test_data = datasets.FashionMNIST(
    root="./data",
    train=False,
    download=True,
    transform=ToTensor()
)
# 标签字典，一个key键对应一个label
labels_map = {
    0: "T-Shirt",
    1: "Trouser",
    2: "Pullover",
    3: "Dress",
    4: "Coat",
    5: "Sandal",
    6: "Shirt",
    7: "Sneaker",
    8: "Bag",
    9: "Ankle Boot",
}

# 设置画布大小
# figure = plt.figure(figsize=(8, 8))
# cols, rows = 3, 3
# for i in range(1, cols * rows + 1):
#     # 随机生成一个索引
#     sample_idx = torch.randint(len(training_data), size=(1,)).item()
#     # 获取样本及其对应的标签
#     img, label = training_data[sample_idx]
#     figure.add_subplot(rows, cols, i)
#     # 设置标题
#     plt.title(labels_map[label])
#     # 不显示坐标轴
#     plt.axis("off")
#     # 显示灰度图
#     plt.imshow(img.squeeze(), cmap="gray")
# plt.show()

# 训练数据加载器; 根据数据集生成一个迭代对象，用于模型的训练
train_dataloader = DataLoader(
    # 定义好的数据集
    dataset=training_data,
    # 设置批量大小
    batch_size=128,
    # 线程数，默认为0。在Windows下设置大于0的数可能会报错。
    num_workers=0,
    # 打乱样本的顺序
    shuffle=True)
# 测试数据加载器
test_dataloader = DataLoader(
    dataset=test_data,
    batch_size=128,
    shuffle=True)

# 展示图片和标签
# train_features, train_labels = next(iter(train_dataloader))
# print(f"Feature batch shape: {train_features.size()}")
# print(f"Labels batch shape: {train_labels.size()}")
# img = train_features[0].squeeze()
# label = train_labels[0]
# plt.imshow(img, cmap="gray")
# plt.show()
# print(f"Label: {label}")

# 模型定义
class Network(nn.Module):
    def __init__(self):
        super(Network, self).__init__()  # 执行父类中的 init 函数
        self.flatten = nn.Flatten()      # 将每个大小为28x28的图像转换为784个像素值的连续数组
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(in_features=28 * 28, out_features=512), # 线性层
            nn.ReLU(),
            nn.Linear(in_features=512, out_features=512),
            nn.ReLU(),
            nn.Linear(in_features=512, out_features=10),
        )

    def forward(self, x):
        x = self.flatten(x)
        logits = self.linear_relu_stack(x)
        return logits

# 优化模型参数
def train_loop(dataloader, model, loss_func, optimizer, device):
    size = len(dataloader.dataset)
    for batch, (X, y) in enumerate(dataloader):
        X = X.to(device)
        y = y.to(device)
        # 前向传播，计算预测值
        pred = model(X)
        # 计算损失
        loss = loss_func(pred, y)
        # 反向传播，优化参数
        optimizer.zero_grad()  # 将模型的梯度归 0
        loss.backward()        # 用来存储每个参数的损失梯度
        optimizer.step()       # 梯度调整完以后，调整参数

        if batch % 100 == 0:
            loss, current = loss.item(), batch * len(X)
            print(f"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]")

# 测试模型性能
def test_loop(dataloader, model, loss_fn, device):
    size = len(dataloader.dataset)
    num_batches = len(dataloader)
    test_loss, correct = 0, 0

    with torch.no_grad():
        for X, y in dataloader:
            X = X.to(device)
            y = y.to(device)
            # 前向传播，计算预测值
            pred = model(X)
            # 计算损失
            test_loss += loss_fn(pred, y).item()
            # 计算准确率
            correct += (pred.argmax(1) == y).type(torch.float).sum().item()

    test_loss /= num_batches
    correct /= size
    print(
        f"Test Error: \n Accuracy: {(100 * correct):>0.1f}%, Avg loss: {test_loss:>8f} \n")

# input
if __name__ == '__main__':
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"Using {device} device")
    # 定义模型，并将模型移动到设备上
    model = Network().to(device)
    # 设置超参数
    learning_rate = 1e-3
    epochs = 20
    # 定义损失函数和优化器
    loss_func = nn.CrossEntropyLoss()
    optimizer = torch.optim.SGD(params=model.parameters(), lr=learning_rate)
    # 训练模型
    for t in range(epochs):
        print(f"Epoch {t + 1}\n-------------------------------")
        train_loop(train_dataloader, model, loss_func, optimizer, device)
        test_loop(test_dataloader, model, loss_func, device)
    print("Done!")
    # 保存模型
    torch.save(model.state_dict(), 'model_weights.pth')