Jittor MNIST 分类教程

完整代码：https://github.com/Jittor/mnistclassification-jittor

MNIST介绍 ：

MNIST 数据集可在 http://yann.lecun.com/exdb/mnist/ 下载, 它是一个对0到9十个数字进行分类的数据集。它包含了四个部分:

训练图像: train-images-idx3-ubyte.gz (9.9 MB, 解压后 47 MB, 包含 60,000 个样本)

训练标签: train-labels-idx1-ubyte.gz (29 KB, 解压后 60 KB, 包含 60,000 个标签)

测试图像: t10k-images-idx3-ubyte.gz (1.6 MB, 解压后 7.8 MB, 包含 10,000 个样本)

测试标签: t10k-labels-idx1-ubyte.gz (5KB, 解压后 10 KB, 包含 10,000 个标签)

MNIST 数据集如下图所示

示例2：MNIST图像分类 - 图1

使用 Jittor 对 MNIST 进行分类

1.首先第一步，需要引入相关的依赖，如下所示。

  1. # classification mnist example 
  2. import jittor as jt  # 将 jittor 引入
  3. from jittor import nn, Module  # 引入相关的模块
  4. import numpy as np
  5. import sys, os
  6. import random
  7. import math 
  8. from jittor import init
  9. jt.flags.use_cuda = 1 # jt.flags.use_cuda 表示是否使用 gpu 训练。
  10. # 如果 jt.flags.use_cuda=1，表示使用GPU训练 如果 jt.flags.use_cuda = 0 表示使用 CPU
  11. from jittor.dataset.mnist import MNIST 
  12. #由于 MNIST 是一个常见的数据集，其数据载入已经被封装进 jittor 所以可以直接调用。
  13. import matplotlib.pyplot as plt
  14. import pylab as pl # 用于绘制 Loss 曲线 和 MNIST 数据

2.加载 MNIST 数据集，需要继承 Dataset 类，需要实现类中的 init() 和 getitem() 函数，对于MNIST，实现如下所示.

  1. from jittor.dataset import Dataset
  3. class MNIST(Dataset):
  4.     def __init__(self, data_root="/mnist_data/", train=True ,download=True, batch_size=1, shuffle=False):
  5.         # if you want to test resnet etc you should set input_channel = 3, because the net set 3 as the input dimensions
  6.         super().__init__()
  7.         self.data_root = data_root
  8.         self.batch_size = batch_size
  9.         self.shuffle = shuffle
  10.         self.is_train = train
  11.         if download == True:
  12.             self.download_url()
  14.         filesname = [
  15.                 "train-images-idx3-ubyte.gz",
  16.                 "t10k-images-idx3-ubyte.gz",
  17.                 "train-labels-idx1-ubyte.gz",
  18.                 "t10k-labels-idx1-ubyte.gz"
  19.         ]
  20.         self.mnist = {}
  21.         if self.is_train:
  22.             with gzip.open(data_root + filesname[0], 'rb') as f:
  23.                 self.mnist["images"] = np.frombuffer(f.read(), np.uint8, offset=16).reshape(-1,28, 28)
  24.             with gzip.open(data_root + filesname[2], 'rb') as f:
  25.                 self.mnist["labels"] = np.frombuffer(f.read(), np.uint8, offset=8)
  26.         else:
  27.             with gzip.open(data_root + filesname[1], 'rb') as f:
  28.                 self.mnist["images"] = np.frombuffer(f.read(), np.uint8, offset=16).reshape(-1,28, 28)
  29.             with gzip.open(data_root + filesname[3], 'rb') as f:
  30.                 self.mnist["labels"] = np.frombuffer(f.read(), np.uint8, offset=8)
  31.         assert(self.mnist["images"].shape[0] == self.mnist["labels"].shape[0])
  32.         self.total_len = self.mnist["images"].shape[0]
  33.         # this function must be called
  34.         self.set_attrs(batch_size = self.batch_size, total_len=self.total_len, shuffle= self.shuffle)
  35.     def __getitem__(self, index):
  36.         img = Image.fromarray (self.mnist['images'][index]) 
  37.         img = np.array (img)
  38.         img = img[np.newaxis, :]
  39.         return np.array((img / 255.0), dtype = np.float32), self.mnist['labels'][index]
  41.     def download_url(self):
  42.         resources = [
  43.             ("http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz", "f68b3c2dcbeaaa9fbdd348bbdeb94873"),
  44.             ("http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz", "d53e105ee54ea40749a09fcbcd1e9432"),
  45.             ("http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz", "9fb629c4189551a2d022fa330f9573f3"),
  46.             ("http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz", "ec29112dd5afa0611ce80d1b7f02629c")
  47.         ]
  49.         for url, md5 in resources:
  50.             filename = url.rpartition('/')[2]
  51.             download_url_to_local(url, filename, self.data_root, md5)

3.模型的定义：我们定义模型需要继承 jittor 的 Module 类。需要实现 init 函数和 execute 函数。init 用于定义模型由哪些操作组成， execute 函数定义了模型执行的顺序和模型的返回值。

  1. class Model (Module):
  2.     def __init__ (self):
  3.         super (Model, self).__init__()
  4.         self.conv1 = nn.Conv (1, 32, 3, 1) # no padding
  6.         self.conv2 = nn.Conv (32, 64, 3, 1)
  7.         self.bn = nn.BatchNorm(64)
  9.         self.max_pool = nn.Pool (2, 2)
  10.         self.relu = nn.Relu()
  11.         self.fc1 = nn.Linear (64 * 12 * 12, 256)
  12.         self.fc2 = nn.Linear (256, 10)
  13.     def execute (self, x) : 
  14.         # it's simliar to forward function in Pytorch 
  15.         x = self.conv1 (x)
  16.         x = self.relu (x)
  18.         x = self.conv2 (x)
  19.         x = self.bn (x)
  20.         x = self.relu (x)
  22.         x = self.max_pool (x)
  23.         x = jt.reshape (x, [x.shape[0], -1])
  24.         x = self.fc1 (x)
  25.         x = self.relu(x)
  26.         x = self.fc2 (x)
  27.         return x

4.对模型进行训练。对模型训练需要定义训练时的超参数，以及需要定义训练过程。训练函数在 train 函数中定义，测试函数在 val 函数中定义。

  1. def train(model, train_loader, optimizer, epoch, losses, losses_idx):
  2.     model.train()
  3.     lens = len(train_loader)
  4.     for batch_idx, (inputs, targets) in enumerate(train_loader):
  5.         outputs = model(inputs)
  6.         loss = nn.cross_entropy_loss(outputs, targets)
  7.         optimizer.step (loss)
  8.         losses.append(loss.data[0])
  9.         losses_idx.append(epoch * lens + batch_idx)
  11.         if batch_idx % 10 == 0:
  12.             print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
  13.                 epoch, batch_idx, len(train_loader) ,
  14.                 100. * batch_idx / len(train_loader), loss.data[0]))
  16. def val(model, val_loader, epoch):
  17.     model.eval()
  19.     test_loss = 0
  20.     correct = 0
  21.     total_acc = 0
  22.     total_num = 0
  23.     for batch_idx, (inputs, targets) in enumerate(val_loader):
  24.         batch_size = inputs.shape[0]
  25.         outputs = model(inputs)
  26.         pred = np.argmax(outputs.data, axis=1)
  27.         acc = np.sum(targets.data==pred)
  28.         total_acc += acc
  29.         total_num += batch_size
  30.         acc = acc / batch_size
  31.         print(f'Test Epoch: {epoch} [{batch_idx}/{len(val_loader)}]\tAcc: {acc:.6f}')        
  32.         print('Test Acc =', total_acc / total_num)
  34. batch_size = 64
  35. learning_rate = 0.1
  36. momentum = 0.9
  37. weight_decay = 1e-4
  38. epochs = 1
  39. losses = []
  40. losses_idx = []
  41. train_loader = MNIST (train=True, batch_size=batch_size, shuffle=True)
  43. val_loader = MNIST (train=False, batch_size=1, shuffle=False)
  45. model = Model ()
  46. optimizer = nn.SGD(model.parameters(), learning_rate, momentum, weight_decay)
  47. for epoch in range(epochs):
  48.     train(model, train_loader, optimizer, epoch, losses, losses_idx)
  49.     val(model, val_loader, epoch)

5.绘制 Loss 曲线 : 将 Loss 曲线进行可视化。

  1. pl.plot(losses_idx, losses)
  2. pl.xlabel('Iterations')
  3. pl.ylabel('Train_loss')

6.存储模型：模型训练完成需要存储起来,下面代码展示了 Jittor 如何进行存储模型。

  1. model_path = './mnist_model.pkl'
  2. model.save(model_path)

7.加载模型并对模型进行测试，下面展示了 Jittor 如何加载模型，并对模型进行测试。

  1. def vis_img(img):
  2.     np_img = img.data.reshape([28, 28])
  3.     plt.imshow(np_img, cmap='gray') 
  6. new_model = Model()
  7. new_model.load_parameters(jt.load(model_path))
  8. data_iter = iter(val_loader)
  9. val_data, val_label = next(data_iter)
  10. print (val_label.shape)
  11. outputs = new_model(val_data)
  12. prediction = np.argmax(outputs.data, axis=1)
  13. print (prediction)
  14. print (val_label)
  16. vis_img(val_data)