pytorch中如何实现手写数字图片识别

小编给大家分享一下pytorch中如何实现手写数字图片识别，希望大家阅读完这篇文章之后都有所收获，下面让我们一起去探讨吧！

具体内容如下

数据集：MNIST数据集，代码中会自动下载，不用自己手动下载。数据集很小，不需要GPU设备，可以很好的体会到pytorch的魅力。
模型+训练+预测程序：

import torchfrom torch import nnfrom torch.nn import functional as Ffrom torch import optimimport torchvisionfrom matplotlib import pyplot as pltfrom utils import plot_image, plot_curve, one_hot# step1  load datasetbatch_size = 512train_loader = torch.utils.data.DataLoader(    torchvision.datasets.MNIST('mnist_data', train=True, download=True,                               transform=torchvision.transforms.Compose([                                   torchvision.transforms.ToTensor(),                                   torchvision.transforms.Normalize(                                       (0.1307,), (0.3081,)                                   )                               ])),    batch_size=batch_size, shuffle=True)test_loader = torch.utils.data.DataLoader(    torchvision.datasets.MNIST('mnist_data/', train=False, download=True,                               transform=torchvision.transforms.Compose([                                   torchvision.transforms.ToTensor(),                                   torchvision.transforms.Normalize(                                       (0.1307,), (0.3081,)                                   )                               ])),    batch_size=batch_size, shuffle=False)x , y = next(iter(train_loader))print(x.shape, y.shape, x.min(), x.max())plot_image(x, y, "image_sample")class Net(nn.Module):    def __init__(self):        super(Net, self).__init__()        self.fc1 = nn.Linear(28*28, 256)        self.fc2 = nn.Linear(256, 64)        self.fc3 = nn.Linear(64, 10)    def forward(self, x):        # x: [b, 1, 28, 28]        # h2 = relu(xw1 + b1)        x = F.relu(self.fc1(x))        # h3 = relu(h2w2 + b2)        x = F.relu(self.fc2(x))        # h4 = h3w3 + b3        x = self.fc3(x)        return xnet = Net()optimizer = optim.SGD(net.parameters(), lr=0.01, momentum=0.9)train_loss = []for epoch in range(3):    for batch_idx, (x, y) in enumerate(train_loader):        #加载进来的图片是一个四维的tensor，x: [b, 1, 28, 28], y:[512]        #但是我们网络的输入要是一个一维向量（也就是二维tensor)，所以要进行展平操作        x = x.view(x.size(0), 28*28)        #  [b, 10]        out = net(x)        y_onehot = one_hot(y)        # loss = mse(out, y_onehot)        loss = F.mse_loss(out, y_onehot)        optimizer.zero_grad()        loss.backward()        # w' = w - lr*grad        optimizer.step()        train_loss.append(loss.item())        if batch_idx % 10 == 0:            print(epoch, batch_idx, loss.item())plot_curve(train_loss)    # we get optimal [w1, b1, w2, b2, w3, b3]total_correct = 0for x,y in test_loader:    x = x.view(x.size(0), 28*28)    out = net(x)    # out: [b, 10]    pred = out.argmax(dim=1)    correct = pred.eq(y).sum().float().item()    total_correct += correcttotal_num = len(test_loader.dataset)acc = total_correct/total_numprint("acc:", acc)x, y = next(iter(test_loader))out = net(x.view(x.size(0), 28*28))pred = out.argmax(dim=1)plot_image(x, pred, "test")

主程序中调用的函数（注意命名为utils)：

import  torchfrom    matplotlib import pyplot as pltdef plot_curve(data):    fig = plt.figure()    plt.plot(range(len(data)), data, color='blue')    plt.legend(['value'], loc='upper right')    plt.xlabel('step')    plt.ylabel('value')    plt.show()def plot_image(img, label, name):    fig = plt.figure()    for i in range(6):        plt.subplot(2, 3, i + 1)        plt.tight_layout()        plt.imshow(img[i][0]*0.3081+0.1307, cmap='gray', interpolation='none')        plt.title("{}: {}".format(name, label[i].item()))        plt.xticks([])        plt.yticks([])    plt.show()def one_hot(label, depth=10):    out = torch.zeros(label.size(0), depth)    idx = torch.LongTensor(label).view(-1, 1)    out.scatter_(dim=1, index=idx, value=1)    return out

看完了这篇文章，相信你对"pytorch中如何实现手写数字图片识别"有了一定的了解，如果想了解更多相关知识，欢迎关注行业资讯频道，感谢各位的阅读！