pytorch实现mnist数据集分类

1.数据准备

  pytorch框架内置的torchvision中的datasets类中有一些常见的数据集。本文采用随机梯度下降的算法来训练数据，数据集构造如下所示：

 # 将数据集转换成张量并且归一化
 transform = transforms.Compose({
     transforms.ToTensor(),
     transforms.Normalize((0.1307,), (0.3081,))}
 )
 # 下载数据集，download= True表示从网络下载，本文已经下载好了这里设置为False
 train_data = datasets.MNIST(root = './data/mnist', train=True, download=False, transform=transform)
 test_data = datasets.MNIST(root='./data/mnist', train=False, download=False, transform=transform)
 # 构造小批量数据，训练集一般要随机，所以shuffle=True
 train_loader = DataLoader(train_data, batch_size=32, shuffle=True)
 # 测试集不用shuffle
 test_loader = DataLoader(test_data, batch_size=32, shuffle=False)

2.模型构造

  本文使用卷积神经网络实现分类，模型的输入为(batch_size,3,28,28)，输出为(batch_size,10)。第一层卷积操作中卷积核大小为5，一个有10个卷积核，池化层大小为2；第二个卷积操作中卷积核大小为5，一个有20个卷积核，池化层大小为2；第三层我全连接层，在进行全连接操作需要将tensor尺寸变为(batch_size, -1)，具体如下所示：

class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = torch.nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = torch.nn.Conv2d(10, 20, kernel_size=5)
        self.pooling = torch.nn.MaxPool2d(2)
        self.fc = torch.nn.Linear(320, 10)

    def forward(self, x):
        batch_size = x.size(0) # batch_size = 32
        x = F.relu(self.pooling(self.conv1(x)))
        x = F.relu(self.pooling(self.conv2(x)))
        x = x.view(batch_size, -1)
        return self.fc(x)

需要注意的是全连接后不需要进行激活操作，交叉熵损失函数里面自动进行了softmax操作

3.训练和测试

  测试的时候不需要计算产生张量图

def train(epoch):

    running_loss = 0
    for i,data in enumerate(train_loader, 0):
        inputs, label = data
        y_predict = model(inputs)
        loss = criterion(y_predict, label)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
        if i % 300 == 299:
            print('[%d, %5d] loss: %.3f' % (epoch + 1, i + 1, running_loss / 300))
            running_loss = 0.0

def test():
    correct = 0
    total = 0
    with torch.no_grad():
        for data in test_loader:
            inputs, label = data
            y_pred = model(inputs)
            _, predicted = torch.max(y_pred.data, dim=1)
            total += label.size(0)
            correct += (predicted == label).sum().item()
    print('Accuracy on test set:%d %%' % (100 * correct / total))

4.模型训练

  本文采用是交叉熵损失函数，随机梯度下降算法

model = Net()
criterion = torch.nn.CrossEntropyLoss(size_average=True)
optimizer = torch.optim.SGD(model.parameters(), lr = 0.01)

for epoch in range(10):
    train(epoch)
    test()

5.完整代码

import torch
from torch.functional import F
from torch.utils.data import DataLoader
from torchvision import datasets, transforms


class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = torch.nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = torch.nn.Conv2d(10, 20, kernel_size=5)
        self.pooling = torch.nn.MaxPool2d(2)
        self.fc = torch.nn.Linear(320, 10)

    def forward(self, x):
        batch_size = x.size(0) # batch_size = 32
        x = F.relu(self.pooling(self.conv1(x)))
        x = F.relu(self.pooling(self.conv2(x)))
        x = x.view(batch_size, -1)
        return self.fc(x)

def train(epoch):

    running_loss = 0
    for i,data in enumerate(train_loader, 0):
        inputs, label = data
        y_predict = model(inputs)
        loss = criterion(y_predict, label)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
        if i % 300 == 299:
            print('[%d, %5d] loss: %.3f' % (epoch + 1, i + 1, running_loss / 300))
            running_loss = 0.0

def test():
    correct = 0
    total = 0
    with torch.no_grad():
        for data in test_loader:
            inputs, label = data
            y_pred = model(inputs)
            _, predicted = torch.max(y_pred.data, dim=1)
            total += label.size(0)
            correct += (predicted == label).sum().item()
    print('Accuracy on test set:%d %%' % (100 * correct / total))


if __name__ == '__main__':
    # 将数据集转换成张量并且归一化
    transform = transforms.Compose({
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))}
    )
    # 下载数据集，download= True表示从网络下载，本文已经下载好了这里设置为False
    train_data = datasets.MNIST(root = './data/mnist', train=True, download=False, transform=transform)
    test_data = datasets.MNIST(root='./data/mnist', train=False, download=False, transform=transform)
    # 构造小批量数据，训练集一般要随机，所以shuffle=True
    train_loader = DataLoader(train_data, batch_size=32, shuffle=True)
    # 测试集不用shuffle
    test_loader = DataLoader(test_data, batch_size=32, shuffle=False)

    model = Net()
    criterion = torch.nn.CrossEntropyLoss(size_average=True)
    optimizer = torch.optim.SGD(model.parameters(), lr = 0.01)

    for epoch in range(10):
        train(epoch)
        test()

有不懂的地方可以加Q:1781983207，欢迎进行交流。