PyTorch深度学习实践——10 卷积神经网络（基础篇）

UP：B站-刘二大人

原视频链接：10.卷积神经网络（基础篇）_哔哩哔哩_bilibili

import torch
from torchvision import datasets
from torchvision import transforms
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt

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 Model(torch.nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.conv1 = torch.nn.Conv2d(1, 10, kernel_size = 5) # 三个参数分别是：输入的通道数，输出的通道数，卷积核的大小（5×5）
        self.conv2 = torch.nn.Conv2d(10, 20, kernel_size = 3, padding=1)
        self.conv3 = torch.nn.Conv2d(20, 32, kernel_size = 3)
        self.pooling = torch.nn.MaxPool2d(2) # 2*2的池化窗口
        self.active = torch.nn.ReLU()
        self.fc1 = torch.nn.Linear(128, 64) # 全连接层
        self.fc2 = torch.nn.Linear(64, 32)
        self.fc3 = torch.nn.Linear(32, 10)

    def forward(self, x):
        x = self.pooling(self.active(self.conv1(x)))
        x = self.pooling(self.active(self.conv2(x)))
        x = self.pooling(self.active(self.conv3(x)))
        # 此时的x是维度为4的tensor，即(batch_size, c , H, W),x.size(0)指的是batch_size的值
        x = x.view(x.size(0), -1)
        x = self.fc1(x)
        x = self.fc2(x)
        x = self.fc3(x)
        return x

model = Model()

criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.5)

loss_list = []
def train(epoch):
    loss_sum = 0.0
    for i, (inputs, labels) in enumerate(train_loader):
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss_sum += loss.item()
        loss_list.append(loss.item())

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

        if i % 300 ==299:
            print("[%d %5d] loss:%.3f" % (epoch+1, i+1, loss_sum / i ))

#if __name__=="__main__":
   #for epoch in range(10):
        #train(epoch)


def test():
    correct = 0
    total = 0
    with torch.no_grad(): # 测试不用计算梯度
        for (images, targets) in test_loader:
            y_pred = model(images)
            _, predicted = torch.max(y_pred.data, dim=1)
            total += targets.size(0)
            correct += (predicted == targets).sum().item()
        print("Accuarcy on test dataset:%.3f %% [%d %d]"%(100*correct/total, correct, total))

if __name__=="__main__":
    for epoch in range(5):
        train(epoch)
        test()

plt.plot(range(len(loss_list)), loss_list)
plt.xlabel("step")
plt.ylabel("loss")
plt.show()