Step 4: Train the neural network

In this step, you learn how to train the previously defined network with data. First, import the libraries. The new ones are mxnet.init for more weight initialization methods. Import the datasets and transforms to load and transform computer vision datasets. Import matplotlib for drawing, and time for benchmarking. The example command here shows this.

# Uncomment the following line if matplotlib is not installed.
# !pip install matplotlib

from mxnet import np, npx, gluon, init, autograd
from mxnet.gluon import nn
from IPython import display
import matplotlib.pyplot as plt
import time
npx.set_np()

Get data

The handwritten digit, MNIST dataset is one of the most commonly used datasets in deep learning. However, it’s too simple to get 99 percent accuracy. For this tutorial, you use a similar but slightly more complicated dataset called FashionMNIST. The end-goal is to classify clothing types.

The dataset can be automatically downloaded through Gluon’s data.vision.datasets module. The following code downloads the training dataset and shows the first example.

mnist_train = gluon.data.vision.datasets.FashionMNIST(train=True)
X, y = mnist_train[0]
('X shape: ', X.shape, 'X dtype', X.dtype, 'y:', y)

('X shape: ',
 (28, 28, 1),
 'X dtype',
 dtype('uint8'),
 'y:',
 array(2, dtype=int32))

Each example in this dataset is a \(28\times 28\) size grey image, which is presented as ndarray with the shape format of (height, width, channel). The label is a numpy scalar.

Next, visualize the first six examples.

text_labels = ['t-shirt', 'trouser', 'pullover', 'dress', 'coat',
               'sandal', 'shirt', 'sneaker', 'bag', 'ankle boot']
X, y = mnist_train[0:10]
# plot images
display.set_matplotlib_formats('svg')
_, figs = plt.subplots(1, X.shape[0], figsize=(15, 15))
for f, x, yi in zip(figs, X, y):
    # 3D->2D by removing the last channel dim
    f.imshow(x.reshape((28,28)).asnumpy())
    ax = f.axes
    ax.set_title(text_labels[int(yi)])
    ax.title.set_fontsize(14)
    ax.get_xaxis().set_visible(False)
    ax.get_yaxis().set_visible(False)
plt.show()

In order to feed data into a Gluon model, convert the images to the (channel, height, width) format with a floating point data type. It can be done by transforms.ToTensor. In addition, normalize all pixel values with transforms.Normalize with the real mean 0.13 and standard deviation 0.31. You can chain these two transforms together and apply it to the first element of the data pair, namely the images.

transformer = gluon.data.vision.transforms.Compose([
    gluon.data.vision.transforms.ToTensor(),
    gluon.data.vision.transforms.Normalize(0.13, 0.31)])
mnist_train = mnist_train.transform_first(transformer)

FashionMNIST is a subclass of gluon.data.Dataset, which defines how to get the i-th example. In order to use it in training, you need to get a (randomized) batch of examples. Do this by using gluon.data.DataLoader. The example here uses four works to process data in parallel, which is often necessary especially for complex data transforms.

batch_size = 256
train_data = gluon.data.DataLoader(
    mnist_train, batch_size=batch_size, shuffle=True, num_workers=4)

The returned train_data is an iterable object that yields batches of images and labels pairs.

for data, label in train_data:
    print(data.shape, label.shape)
    break

(256, 1, 28, 28) (256,)

Finally, create a validation dataset and data loader.

mnist_valid = gluon.data.vision.FashionMNIST(train=False)
valid_data = gluon.data.DataLoader(
    mnist_valid.transform_first(transformer),
    batch_size=batch_size, num_workers=4)

Define the model

Implement the network called LeNet. One difference here is that you change the weight initialization method to Xavier, which is a popular choice for deep convolutional neural networks.

net = nn.Sequential()
net.add(nn.Conv2D(channels=6, kernel_size=5, activation='relu'),
        nn.MaxPool2D(pool_size=2, strides=2),
        nn.Conv2D(channels=16, kernel_size=3, activation='relu'),
        nn.MaxPool2D(pool_size=2, strides=2),
        nn.Dense(120, activation="relu"),
        nn.Dense(84, activation="relu"),
        nn.Dense(10))
net.initialize(init=init.Xavier())

In addition to the neural network, define the loss function and optimization method for training. Use standard softmax cross entropy loss for classification problems. It first performs softmax on the output to obtain the predicted probability, and then compares the label with the cross entropy.

softmax_cross_entropy = gluon.loss.SoftmaxCrossEntropyLoss()

The optimization method you pick is the standard stochastic gradient descent with constant learning rate of 0.1.

trainer = gluon.Trainer(net.collect_params(), 'sgd', {'learning_rate': 0.1})

The trainer is created with all parameters (both weights and gradients) in net. Later on, you only need to call the step method to update its weights.

Train the model

Create an auxiliary function to calculate the model accuracy.

def acc(output, label):
    # output: (batch, num_output) float32 ndarray
    # label: (batch, ) int32 ndarray
    return (output.argmax(axis=1) == label.astype('float32')).mean()

Implement the complete training loop.

for epoch in range(10):
    train_loss, train_acc, valid_acc = 0., 0., 0.
    tic = time.time()
    for data, label in train_data:
        # forward + backward
        with autograd.record():
            output = net(data)
            loss = softmax_cross_entropy(output, label)
        loss.backward()
        # update parameters
        trainer.step(batch_size)
        # calculate training metrics
        train_loss += loss.mean()
        train_acc += acc(output, label)
    # calculate validation accuracy
    for data, label in valid_data:
        valid_acc += acc(net(data), label)
    print("Epoch %d: loss %.3f, train acc %.3f, test acc %.3f, in %.1f sec" % (
            epoch, train_loss/len(train_data), train_acc/len(train_data),
            valid_acc/len(valid_data), time.time()-tic))

Epoch 0: loss 0.843, train acc 0.683, test acc 0.800, in 5.9 sec
Epoch 1: loss 0.494, train acc 0.817, test acc 0.850, in 3.1 sec
Epoch 2: loss 0.422, train acc 0.844, test acc 0.850, in 2.9 sec
Epoch 3: loss 0.380, train acc 0.861, test acc 0.863, in 2.9 sec
Epoch 4: loss 0.358, train acc 0.869, test acc 0.872, in 2.9 sec
Epoch 5: loss 0.336, train acc 0.875, test acc 0.876, in 3.0 sec
Epoch 6: loss 0.321, train acc 0.880, test acc 0.876, in 2.9 sec
Epoch 7: loss 0.309, train acc 0.885, test acc 0.881, in 3.0 sec
Epoch 8: loss 0.299, train acc 0.889, test acc 0.888, in 2.9 sec
Epoch 9: loss 0.292, train acc 0.891, test acc 0.886, in 3.0 sec

Save the model

Finally, save the trained parameters onto disk, so that you can use them later.

net.save_parameters('net.params')

Next Steps

After the model is trained and saved, learn how to use it to predict new examples: Step 5: Predict with a pretrained model.