写在最前边

这篇文章要写的内容看封面，就是要用一篇文章讲解一下，怎么用Fashion-MNIST数据集，我们自己建一个神经网络，训练好之后用它做图片分类。

import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensor

老规矩，一切开始之前我们先浅浅导个包。

数据集

Pytorch梳理数据集有两个主要的方法： torch.utils.data.DataLoader和 torch.utils.data.Dataset。

Dataset：存储数据集的样本及其标签
DataLoader 将数据集包裹到迭代器中，方便模型读取

Pytorch提供了人工智能三大领域（文本、图像、语音）主要的数据集，使用TorchText, TorchVision, TorchAudio我们可以拿到对应的一些主流数据集。

今天我们要做一个图片分类网络，所以我们在这里用TorchVision。

我们使用如下代码将FashionMNIST的训练集和测试集依次加载进来。

# 从FashionMNIST加载训练集
training_data = datasets.FashionMNIST(
    root="data",
    train=True,
    download=True,
    transform=ToTensor(),
)

# 从FashionMNIST加载测试集
test_data = datasets.FashionMNIST(
    root="data",
    train=False,
    download=True,
    transform=ToTensor(),
)

batch_size = 64

# 使用DataLoader创建好迭代器加载数据，方便模型读取
train_dataloader = DataLoader(training_data, batch_size=batch_size)
test_dataloader = DataLoader(test_data, batch_size=batch_size)

for X, y in test_dataloader:
    print(f"Shape of X [N, C, H, W]: {X.shape}")
    print(f"Shape of y: {y.shape} {y.dtype}")
    break

我们用for循环读了一步测试集DataLoader的内容，可以看到如下输出：

Shape of X [N, C, H, W]: torch.Size([64, 1, 28, 28]) Shape of y: torch.Size([64]) torch.int64

X是样本，我们可以看到样本的张量形状为[64, 1, 28, 28]，意思是一个batch有64张图，单通道（黑白图），分辨率为28×28.

y是标签，FashionMNIST是分十类的，标签使用整数表示的，可以看到这里y是一个长为64的向量。

创建模型

我们创建自定义模型要作为nn.Module的子类，在__inti__中写网络结构，在forward中写计算图前馈计算的过程。如果有GPU我们就将代码丢到GPU上运行。

# 看看你有无GPU，有GPU就用GPU，没有就用CPU。
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using {device} device")

# 定义模型结构
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 logits

model = NeuralNetwork().to(device)
print(model)

# 这里我是有GPU的，所以显示使用cuda Using cuda device NeuralNetwork( $\quad$ $\quad$ (flatten): Flatten(start_dim=1, end_dim=-1) $\quad$ $\quad$ (linear_relu_stack): $\quad$ $\quad$ Sequential( $\quad$ $\quad$ $\quad$ (0): Linear(in_features=784, out_features=512, bias=True) $\quad$ $\quad$ $\quad$ (1): ReLU() $\quad$ $\quad$ $\quad$ (2): Linear(in_features=512, out_features=512, bias=True) $\quad$ $\quad$ $\quad$ (3): ReLU() $\quad$ $\quad$ $\quad$ (4): Linear(in_features=512, out_features=10, bias=True) $\quad$ $\quad$ $\quad$ ) $\quad$ $\quad$)

优化网络

训练网络之前先定义好loss优化目标和优化器。

loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)

训练网络我们要很多epochs，在每次循环中，模型都会使用训练集预测结果，使用预测结果和真实结果进行比较，通过反向传播不断修正模型的参数。

def train(dataloader, model, loss_fn, optimizer):
    size = len(dataloader.dataset)
    model.train()
    for batch, (X, y) in enumerate(dataloader):
        X, y = X.to(device), y.to(device)

        # 使用loss算预测值和真实值的误差
        pred = model(X)
        loss = loss_fn(pred, y)

        # 反向传播更新网络
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # 每100个batch打印一下当前训练进度
        if batch % 100 == 0:
            loss, current = loss.item(), batch * len(X)
            print(f"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]")

我们也要不断检查模型的表现，以便知道他是在往正确的方向上学习的。

def test(dataloader, model, loss_fn):
    size = len(dataloader.dataset)
    num_batches = len(dataloader)
    model.eval（) # 将模型转化为评价模式
    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()
            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")

训练过程就是跑多个循环，我们这里设置epochs=15就是进行15次参数优化。每个epoch中模型都会更新参数以让自己的预测结果更为准确。我们在这里是每轮epoch都打印一下当前的预测的准确率和loss。loss越低，准确率越高，模型越好。

epochs = 15
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!")

我先是跑了5个epoch，感觉准确率有点低，所以加到15了。当然这里是简单数据集，训练过多可能出现过拟合的现象，不能盲目增大epoch提高准确率。

epochs=5
epoch=15

模型训好了，做个预测

classes = [
    "T-shirt/top",
    "Trouser",
    "Pullover",
    "Dress",
    "Coat",
    "Sandal",
    "Shirt",
    "Sneaker",
    "Bag",
    "Ankle boot",
]

# 将模型改为评价模式
model.eval（)
x, y = test_data[0][0], torch.tensor(test_data[0][1])

# 把x，y挪到GPU上
x = x.to(device)
y = y.to(device)

with torch.no_grad():
    pred = model(x)
    predicted, actual = classes[pred[0].argmax(0)], classes[y]
    print(f'Predicted: "{predicted}", Actual: "{actual}"')

关于x和y的取法x, y = test_data[0][0], torch.tensor(test_data[0][1])，test_data长下图这样，test_data[0]就是第一个样本，test_data[0][0]是取样本，test_data[0][1]是取标签。

用一个图像分类实例拿捏Pytorch使用方法

至于.to(device)，x，y都是在CPU上的，但是模型是在GPU上的，我们要将其挪到GPU上，否则会报错RuntimeError:Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

模型保存和加载

模型训练好了，以后可以直接拿来用的。所以我们要将其保存一下：

torch.save(model.state_dict(), "model.pth")
print("Saved PyTorch Model State to model.pth")

Saved PyTorch Model State to model.pth

需要的时候重新加载模型：

model = NeuralNetwork()
model.load_state_dict(torch.load("model.pth"))

因为我们前边建好模型，我们是接着前边代码运行的，所以这里不需要加载模型。

使用模型预测

这里是把类别都放进来了，给模型一个样本，看看它的结果。

classes = [
    "T-shirt/top",
    "Trouser",
    "Pullover",
    "Dress",
    "Coat",
    "Sandal",
    "Shirt",
    "Sneaker",
    "Bag",
    "Ankle boot",
]

model.eval（)
x, y = test_data[0][0], test_data[0][1]

with torch.no_grad():
    pred = model(x)
    predicted, actual = classes[pred[0].argmax(0)], classes[y]
    print(f'Predicted: "{predicted}", Actual: "{actual}"')

Predicted: "Ankle boot", Actual: "Ankle boot"

可以看到在这里是预测对了的。

在这我们不需要.to(device)了，因为这里是使用前边建好的模型+我们自己读进来的参数，这是一个新的网络，我们没有把这个新网络放到GPU上。也就是说你现在这个代码里有两个一模一样的模型。

你在GPU上训练好了这个模型，然后将参数保存到了本地。
你又从本地读进来参数，在CPU上重现了这个模型。

秒客网