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
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
import torch
from torchvision import transforms
from torchvision import datasets
from torch.utils.data import DataLoader
import torch.nn.functional as F
import torch.optim as optim

batch_size = 64
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])
train_dataset = datasets.MNIST(root='./dataset/mnist/',
                               train=True,
                               download=True,
                               transform=transform)
train_loader = DataLoader(dataset=train_dataset,
                          shuffle=True,
                          batch_size=batch_size)
test_dataset = datasets.MNIST(root='./dataset/mnist/',
                              train=False,
                              download=True,
                              transform=transform)
test_loader = DataLoader(dataset=test_dataset,
                         shuffle=False,
                         batch_size=batch_size)

class Net(torch.nn.Module):  # 继承Module类
    def __init__(self):
        super(Net, self).__init__()  # 继承父类
        self.conv1 = torch.nn.Conv2d(1, 10, kernel_size=5)  # 卷积层1,使用2维的卷积核,1通道->10通道,卷积核大小为5x5,
        self.pooling = torch.nn.MaxPool2d(2)  # 池化层,使用二维的最大池化层,图像大小的height和width都减半并向下取整
        self.conv2 = torch.nn.Conv2d(10, 20, kernel_size=5)  # 卷积层2, 10通道->20通道,卷积核大小为5x5,
        # (batch_size,13,13,10)->(batch_size,11,11,20)
        self.l1 = torch.nn.Linear(320, 256)  # 全连接层
        self.l2 = torch.nn.Linear(256, 128)
        self.l3 = torch.nn.Linear(128, 64)
        self.l4 = torch.nn.Linear(64, 10)

    def forward(self, x):
        batch_size = x.size(0)
        x = F.gelu(self.pooling(self.conv1(
            x)))  # (batch_size,28,28,1)->(batch_size,26,26,10)->(batch_size,13,13,10) 因为没有做填充(padding)所以图像的大小减小了
        x = F.gelu(self.pooling(self.conv2(x)))  # (batch_size,13,13,10)->(batch_size,6,6,20)->(batch_size,4,4,20)
        x = x.view(batch_size, -1)  # (batch_size,4*4*20)将张量展平接入全连接层
        x = F.gelu(self.l1(x))
        x = F.gelu(self.l2(x))
        x = F.gelu(self.l3(x))
        x = self.l4(x)

        return x  # 通过带有激活函数的线性层获得输出

model = Net()  # 实例化神经网络
device = torch.device('cuda:0' if torch.cuda.is_available() else 'CPU')  # 如果能使用GPU,就将模型移植到GPU上用cuda核心进行并行计算
model.to(device)
criterion = torch.nn.CrossEntropyLoss().to(device)
optimizer = optim.Adam(model.parameters(), lr=0.01)

def train(epoch):
    running_loss = 0
    for batch_idx, data in enumerate(train_loader):
        inputs, target = data
        inputs, target = inputs.to(device), target.to(device)
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, target)
        loss.backward()
        optimizer.step()
        running_loss += loss
        if batch_idx % 300 == 299:
            print('[%d,%5d] loss:%.3f' % (epoch + 1, batch_idx + 1, running_loss / 300))
            running_loss = 0

def test():
    corrent = 0
    total = 0
    with torch.no_grad():
        for data in test_loader:
            images, labels = data
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            _, prediction = torch.max(outputs, dim=1)
            total += labels.size(0)
            corrent += (prediction == labels).sum().item()
    print('Accuracy on test set %.2f %%' % (100 * corrent / total))

if __name__ == '__main__':
    for epoch in range(10):
        train(epoch)
        test()

运行结果: