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')
巴特西

pytorch 入门

最新文章

热门文章
