2021-07-21

手写字体识别代码

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torchvision
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 6, 5)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 4 * 4, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)
        x = x.view(-1, self.num_flat_features(x))
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

    def num_flat_features(self, x):
        size = x.size()[1:]  
        num_features = 1
        for s in size:
            num_features *= s
        return num_features


net = Net()
# create your optimizer
optimizer = optim.SGD(net.parameters(), lr=0.01)

#########################################
n_epochs = 3
batch_size_train = 6
batch_size_test = 10
learning_rate = 0.01
momentum = 0.5
log_interval = 10
random_seed = 1
torch.manual_seed(random_seed)

train_loader = torch.utils.data.DataLoader(
  torchvision.datasets.MNIST('./data/', train=True, download=True,
                             transform=torchvision.transforms.Compose([
                               torchvision.transforms.ToTensor(),
                               torchvision.transforms.Normalize(
                                 (0.1307,), (0.3081,))
                             ])),
  batch_size=batch_size_train, shuffle=True)
test_loader = torch.utils.data.DataLoader(
  torchvision.datasets.MNIST('./data/', train=False, download=True,
                             transform=torchvision.transforms.Compose([
                               torchvision.transforms.ToTensor(),
                               torchvision.transforms.Normalize(
                                 (0.1307,), (0.3081,))
                             ])),
  batch_size=batch_size_test, shuffle=True)
#
# examples = enumerate(test_loader)
# batch_idx, (example_data, example_targets) = next(examples)
# print(example_targets)
# print(example_data.shape)
##################################################
train_losses = []
train_counter = []
test_losses = []
test_counter = [i*len(train_loader.dataset) for i in range(n_epochs + 1)]

############################################################

def train(epoch):

  net.train()
  net.zero_grad()
  for batch_idx, (data, target) in enumerate(train_loader):
    optimizer.zero_grad()
    output = net(data)
    output=F.softmax(output)
    loss = F.nll_loss(output, target)
    loss.backward()
    optimizer.step()
    if batch_idx % log_interval == 0:
      print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
        epoch, batch_idx * len(data), len(train_loader.dataset),
        100. * batch_idx / len(train_loader), loss.item()))

      train_losses.append(loss.item())
      train_counter.append(
        (batch_idx*4) + ((epoch-1)*len(train_loader.dataset)))
      torch.save(net.state_dict(), './model.pth')
      torch.save(optimizer.state_dict(), './optimizer.pth')
#####################
# for i in range(10):
#      train(1)

for batch_idx, (test_data, target) in enumerate(test_loader):
    model_para = torch.load("./model.pth")
    net_test = Net()
    net_test.load_state_dict(model_para)
    prediction = F.softmax(net_test(test_data))
    for i in range(10):
        plt.subplot(2,5,i+1)
        plt.tight_layout()
        plt.imshow(test_data[i][0], cmap='gray', interpolation='none')
        print(prediction[i])
        plt.title("prediction: {}".format(torch.max(prediction[i],0)))
        plt.xticks([])
        plt.yticks([])
        plt.show()