PyTorch入门之【MLP】_pytorch mlp

参考：https://www.bilibili.com/video/BV1uh411A7jJ/?spm_id_from=333.999.0.0&vd_source=98d31d5c9db8c0021988f2c2c25a9620
这里用实例来介绍，用MLP来对数据集MNIST进行训练和测试。

训练

总代码：

import torch
from torchvision.transforms import ToTensor
from torchvision import datasets
import torch.nn as nn

# load MNIST dataset
training_data = datasets.MNIST(
    root='../02_dataset/data',
    train=True,
    download=True,
    transform=ToTensor()
)

train_data_loader = torch.utils.data.DataLoader(training_data, batch_size=64, shuffle=True)

# define a MLP model
class MLP(nn.Module):
    def __init__(self):
        super(MLP, self).__init__()
        self.flatten = nn.Flatten()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(28*28, 64), # 28*28 is the size of the image
            nn.ReLU(),
            nn.BatchNorm1d(64),
            nn.Linear(64, 10), # 10 is the number of classes
        )

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

# create a MLP model
device = 'cuda' if torch.cuda.is_available() else 'cpu'
mlp = MLP().to(device)
optimizer = torch.optim.Adam(mlp.parameters(), lr=1e-3)
loss_fn = nn.CrossEntropyLoss()

# train the model
num_epochs = 20

for epoch in range(num_epochs):
    print(f'Epoch {epoch+1}\n-------------------------------')
    for idx, (img, label) in enumerate(train_data_loader):
        size = len(train_data_loader.dataset)
        img, label = img.to(device), label.to(device)

        # compute prediction error
        pred = mlp(img)
        loss = loss_fn(pred, label)

        # backpropagation
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if idx % 400 == 0:
            loss, current = loss.item(), idx*len(img)
            print(f'loss: {loss:>7f} [{current:>5d}/{size:>5d}]')

# save the model
torch.save(mlp.state_dict(), 'mlp.pth')
print('Saved PyTorch Model State to mlp.pth')
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63

下面挨个看各个模块：

# load MNIST dataset
training_data = datasets.MNIST(
    root='../02_dataset/data',
    train=True,
    download=True,
    transform=ToTensor()
)

train_data_loader = torch.utils.data.DataLoader(training_data, batch_size=64, shuffle=True)
1
2
3
4
5
6
7
8
9

上述代码就是训练数据集的加载

# define a MLP model
class MLP(nn.Module):
    def __init__(self):
        super(MLP, self).__init__()
        self.flatten = nn.Flatten()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(28*28, 64), # 28*28 is the size of the image
            nn.ReLU(),
            nn.BatchNorm1d(64),
            nn.Linear(64, 10), # 10 is the number of classes
        )

    def forward(self, x):#前向传播函数
        x = self.flatten(x)#用于将输入的图像数据展平为一维向量。
        logits = self.linear_relu_stack(x)
        return logits
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16

上述代码是定义MLP模块。

# create a MLP model
device = 'cuda' if torch.cuda.is_available() else 'cpu'
mlp = MLP().to(device)
optimizer = torch.optim.Adam(mlp.parameters(), lr=1e-3)
loss_fn = nn.CrossEntropyLoss()
1
2
3
4
5

上述代码创建了一个 MLP（多层感知机）模型，并将其移动到适当的计算设备上。
创建一个 Adam 优化器，传入 MLP 模型的参数 mlp.parameters() 和学习率 lr=1e-3。Adam 优化器是一种常用的优化算法，它可以自动调整学习率，以便更好地更新模型参数。

# train the model
num_epochs = 20

for epoch in range(num_epochs):
    print(f'Epoch {epoch+1}\n-------------------------------')
    for idx, (img, label) in enumerate(train_data_loader):
        size = len(train_data_loader.dataset)
        img, label = img.to(device), label.to(device)#将输入图像img和对应标签label移动到正确的计算设备上。

        # compute prediction error
        pred = mlp(img)
        loss = loss_fn(pred, label)

        # backpropagation
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if idx % 400 == 0:
            loss, current = loss.item(), idx*len(img)
            print(f'loss: {loss:>7f} [{current:>5d}/{size:>5d}]')
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21

上述代码就是训练MLP模型。
反向传播和参数更新：

• 首先需要使用optimizer.zero_grad()将模型参数的梯度清零，以确保每个 batch 的梯度计算不会受到上一次 batch 的影响。
• 接下来，利用 loss.backward() 方法计算模型输出与真实标签之间的误差，并自动进行反向传播。这样就能够通过链式法则计算出每个参数对于损失函数的梯度。
• 最后，通过 optimizer.step() 方法根据计算得到的梯度更新模型的参数。优化器会根据定义的优化算法和学习率等参数来更新模型参数，以减小损失函数的值。

通过反向传播和参数更新的过程，模型会不断调整参数，使得预测值更加接近真实标签，从而提高模型的准确性和性能。

# save the model
torch.save(mlp.state_dict(), 'mlp.pth')
print('Saved PyTorch Model State to mlp.pth')
1
2
3

上述代码就是将训练好的模型的参数保存到指定的文件路径

测试

总的代码如下：

import torch
from torchvision import datasets
from torchvision import transforms
from torchvision.transforms import ToTensor
from torchvision.datasets import ImageFolder
import torch.nn as nn

# load test data
test_data = datasets.MNIST(
    root='../02_dataset/data',
    train=False,
    download=True,
    transform=ToTensor()
)
test_data_loader = torch.utils.data.DataLoader(test_data, batch_size=64, shuffle=True)

transform = transforms.Compose([
    transforms.Grayscale(),
    transforms.RandomRotation(10),
    transforms.ToTensor()
])
my_mnist = ImageFolder(root='../02_dataset/my-mnist', transform=transform)
my_mnist_loader = torch.utils.data.DataLoader(my_mnist, batch_size=64, shuffle=True)

# define a MLP model
class MLP(nn.Module):
    def __init__(self):
        super(MLP, self).__init__()
        self.flatten = nn.Flatten()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(28*28, 64), # 28*28 is the size of the image
            nn.ReLU(),
            nn.BatchNorm1d(64),
            nn.Linear(64, 10), # 10 is the number of classes
        )

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

# load the pretrained model
device = 'cuda' if torch.cuda.is_available() else 'cpu'
mlp = MLP()
mlp.load_state_dict(torch.load('mlp.pth', map_location=device))
mlp.eval().to(device)

# test the pretrained model on MNIST test data
size = len(test_data_loader.dataset)
correct = 0

with torch.no_grad():
    for img, label in test_data_loader:
        img, label = img.to(device), label.to(device)
        pred = mlp(img)

        correct += (pred.argmax(1) == label).type(torch.float).sum().item()

correct /= size
print(f'Accuracy on MNIST: {(100*correct):>0.1f}%')

# test the pretrained model on my MNIST test data
size = len(my_mnist_loader.dataset)
correct = 0

with torch.no_grad():
    for img, label in my_mnist_loader:
        img, label = img.to(device), label.to(device)
        pred = mlp(img)

        correct += (pred.argmax(1) == label).type(torch.float).sum().item()

correct /= size
print(f'Accuracy on my MNIST: {(100*correct):>0.1f}%')
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74

接下来看各个模块：

# load test data
test_data = datasets.MNIST(
    root='../02_dataset/data',
    train=False,
    download=True,
    transform=ToTensor()
)
test_data_loader = torch.utils.data.DataLoader(test_data, batch_size=64, shuffle=True)

transform = transforms.Compose([
    transforms.Grayscale(),
    transforms.RandomRotation(10),
    transforms.ToTensor()
])
my_mnist = ImageFolder(root='../02_dataset/my-mnist', transform=transform)
my_mnist_loader = torch.utils.data.DataLoader(my_mnist, batch_size=64, shuffle=True)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16

上述代码是加载数据集自带的测试数据和自己的测试数据

# define a MLP model
class MLP(nn.Module):
    def __init__(self):
        super(MLP, self).__init__()
        self.flatten = nn.Flatten()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(28*28, 64), # 28*28 is the size of the image
            nn.ReLU(),
            nn.BatchNorm1d(64),
            nn.Linear(64, 10), # 10 is the number of classes
        )

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

# load the pretrained model
device = 'cuda' if torch.cuda.is_available() else 'cpu'
mlp = MLP()
mlp.load_state_dict(torch.load('mlp.pth', map_location=device))
mlp.eval().to(device)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22

上述代码是MLP的定义和初始化一个实例。

# test the pretrained model on MNIST test data
size = len(test_data_loader.dataset)
correct = 0
1
2
3

上述代码用来获取测试数据集的大小，correct用来存分类正确的个数

with torch.no_grad():#是一个上下文管理器，用于在进行模型推理时禁用梯度计算，也就是说禁止反向传播。
    for img, label in test_data_loader:
        img, label = img.to(device), label.to(device)
        pred = mlp(img)

        correct += (pred.argmax(1) == label).type(torch.float).sum().item()
        #pred形状为 [batch_size, num_classes]，其中 batch_size 是当前批次的样本数量，num_classes 是类别的数量。
        #.argmax(1)求每一行的最大值的下标，也就是说求0-9中预测概率求的下标
        #pred.argmax(1) == label判断是不是正确的下标
        #.type(torch.float).sum()就是转化成浮点数并求和，因为这里是按一批64个图片计算的故得求和
        #.item()将张量变成标量，就是将其转化成一个数字

correct /= size
print(f'Accuracy on MNIST: {(100*correct):>0.1f}%')
1
2
3
4
5
6
7
8
9
10
11
12
13
14

with torch.no_grad():
    for img, label in my_mnist_loader:
        img, label = img.to(device), label.to(device)
        pred = mlp(img)

        correct += (pred.argmax(1) == label).type(torch.float).sum().item()

correct /= size
print(f'Accuracy on my MNIST: {(100*correct):>0.1f}%')
1
2
3
4
5
6
7
8
9

这个代码和上面的几乎一模一样就不再赘述了。