RNN(LSTM)实现逻辑回归对FashionMNIST数据集进行分类（使用GPU）

2024-02-02 08:41:33

在本文中，我们将使用PyTorch实现递归神经网络（RNN），具体来说是长短期记忆（LSTM）网络，以对FashionMNIST数据集执行逻辑回归分类任务。此外，我们还将利用图形处理单元（GPU）来加速训练过程。

准备工作

首先，我们需要导入必要的库：

import torch as t
import torchvision as tv
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from torch.nn.utils.rnn import pad_sequence
from torch.nn import functional as F
import time
import sys
import os
import math
import numpy as np

接下来，我们加载FashionMNIST数据集并对其进行预处理：

# 数据集的根目录
root = ".fashionmnist"

# 下载数据集
tv.datasets.FashionMNIST(
    root=root,
    train=True,
    download=True,
)
tv.datasets.FashionMNIST(
    root=root,
    train=False,
    download=True,
)

# 预处理数据集
transform = transforms.Compose([transforms.ToTensor()])

# 加载数据集并使用数据加载器
train_data = tv.datasets.FashionMNIST(root=root, train=True, transform=transform)
test_data = tv.datasets.FashionMNIST(root=root, train=False, transform=transform)
train_loader = DataLoader(train_data, batch_size=64, shuffle=True)
test_loader = DataLoader(test_data, batch_size=64, shuffle=False)

定义LSTM网络

现在，我们可以定义我们的LSTM网络：

class LSTM(t.nn.Module):
    def __init__(self, input_dim, hidden_dim, num_layers):
        super(LSTM, self).__init__()

        # LSTM层
        self.lstm = t.nn.LSTM(input_dim, hidden_dim, num_layers, batch_first=True)

        # 全连接层
        self.fc = t.nn.Linear(hidden_dim, 10)

    def forward(self, x):
        # LSTM层
        x, _ = self.lstm(x)

        # 取最后一步的输出
        x = x[:, -1, :]

        # 全连接层
        x = self.fc(x)

        return x

训练模型

接下来，我们需要训练我们的模型：

# 模型实例化
model = LSTM(28, 64, 2)

# 定义损失函数和优化器
criterion = t.nn.CrossEntropyLoss()
optimizer = t.optim.Adam(model.parameters(), lr=0.001)

# 将模型移动到GPU
if t.cuda.is_available():
    model = model.cuda()

# 训练循环
for epoch in range(10):
    for i, (inputs, labels) in enumerate(train_loader):
        if t.cuda.is_available():
            inputs, labels = inputs.cuda(), labels.cuda()

        # 前向传播
        outputs = model(inputs)

        # 计算损失
        loss = criterion(outputs, labels)

        # 反向传播
        loss.backward()

        # 更新权重
        optimizer.step()

        # 清除梯度
        optimizer.zero_grad()

        # 打印训练信息
        if (i + 1) % 100 == 0:
            print("Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}".format(epoch + 1, 10, i + 1, len(train_loader), loss.item()))

# 保存模型
t.save(model.state_dict(), "lstm.pt")

评估模型

最后，让我们评估一下我们的模型：

# 加载模型
model = LSTM(28, 64, 2)
model.load_state_dict(t.load("lstm.pt"))

# 将模型移动到GPU
if t.cuda.is_available():
    model = model.cuda()

# 评估模型
model.eval()
with t.no_grad():
    correct = 0
    total = 0
    for inputs, labels in test_loader:
        if t.cuda.is_available():
            inputs, labels = inputs.cuda(), labels.cuda()

        # 前向传播
        outputs = model(inputs)

        # 计算准确率
        _, predicted = t.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

    print("Accuracy of the network on the 10000 test images: {} %".format(100 * correct / total))