The MNIST dataset contains 60,000 training and 10,000 test grayscale images of handwritten digits, each 28×28 pixels. A complete MindSpore workflow loads this data, defines a feed-forward neural network, optimizes the parameters, and persists the trained weights.

Data Loading

Place the raw files under MNIST_Data/ with train/ and test/ subdirectories. Use MnistDataset to stream samples and batch them for training.

import mindspore.dataset as ds

train_dataset = ds.MnistDataset(dataset_dir="MNIST_Data/train", shuffle=True)
train_dataset = train_dataset.batch(64)

Network Architecture

Custom models extend nn.Cell. Define submodules inside __init__ and describe the tensor flow inside construct.

import mindspore.nn as nn

class DigitNet(nn.Cell):
    def __init__(self):
        super().__init__()
        self.flatten = nn.Flatten()
        self.backbone = nn.SequentialCell([
            nn.Dense(784, 256),
            nn.ReLU(),
            nn.Dense(256, 128),
            nn.ReLU()
        ])
        self.head = nn.Dense(128, 10)

    def construct(self, inputs):
        x = self.flatten(inputs)
        x = self.backbone(x)
        return self.head(x)

model = DigitNet()

Optimization Loop

Each training step executes forward propagation to produce logits, computes the loss against true labels, derives gradients via automatic differentiation, and updates parameters through the optmiizer.

import mindspore as ms

loss_fn = nn.CrossEntropyLoss(sparse=True, reduction="mean")
optimizer = nn.Adam(model.trainable_params(), learning_rate=1e-3)

class LossWrapper(nn.Cell):
    def __init__(self, network, criterion):
        super().__init__()
        self.network = network
        self.criterion = criterion

    def construct(self, data, label):
        preds = self.network(data)
        return self.criterion(preds, label)

train_net = nn.TrainOneStepCell(LossWrapper(model, loss_fn), optimizer)

for epoch in range(5):
    for sample in train_dataset:
        step_loss = train_net(sample[0], sample[1])

Checkpoint Management

Save the trained parameter values to a checkpoint file.

ms.save_checkpoint(model, "digit_classifier.ckpt")

Resume by creating a new network instance and loading the stored parameters into it.

restored_model = DigitNet()
checkpoint = ms.load_checkpoint("digit_classifier.ckpt")
ms.load_param_into_net(restored_model, checkpoint)