.. _crash_course_train: 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. .. code:: python # 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. .. code:: python mnist_train = gluon.data.vision.datasets.FashionMNIST(train=True) X, y = mnist_train[0] ('X shape: ', X.shape, 'X dtype', X.dtype, 'y:', y) .. parsed-literal:: :class: output ('X shape: ', (28, 28, 1), 'X dtype', dtype('uint8'), 'y:', array(2, dtype=int32)) Each example in this dataset is a :math:`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. .. code:: python 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() .. image:: output_4-train_549aa4_5_0.svg 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. .. code:: python 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. .. code:: python 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. .. code:: python for data, label in train_data: print(data.shape, label.shape) break .. parsed-literal:: :class: output (256, 1, 28, 28) (256,) Finally, create a validation dataset and data loader. .. code:: python 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. .. code:: python 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. .. code:: python softmax_cross_entropy = gluon.loss.SoftmaxCrossEntropyLoss() The optimization method you pick is the standard stochastic gradient descent with constant learning rate of 0.1. .. code:: python 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. .. code:: python 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. .. code:: python 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)) .. parsed-literal:: :class: output 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. .. code:: python net.save_parameters('net.params') Next Steps ---------- After the model is trained and saved, learn how to use it to predict new examples: :ref:`crash_course_predict`.