基于MindSpore的Mnist手写数字分类

发现一个华为出的深度学习框架，同时支持静态图和动态图，尝试一下，将代码收藏起来方便以后使用。

需要先下载数据集:

mkdir -p ./data/train ./data/test
wget -NP ./data/train https://mindspore-website.obs.myhuaweicloud.com/notebook/datasets/mnist/train-labels-idx1-ubyte --no-check-certificate
wget -NP ./data/train https://mindspore-website.obs.myhuaweicloud.com/notebook/datasets/mnist/train-images-idx3-ubyte --no-check-certificate
wget -NP ./data/test https://mindspore-website.obs.myhuaweicloud.com/notebook/datasets/mnist/t10k-labels-idx1-ubyte --no-check-certificate
wget -NP ./data/test https://mindspore-website.obs.myhuaweicloud.com/notebook/datasets/mnist/t10k-images-idx3-ubyte --no-check-certificate

代码如下:

import mindspore
from mindspore import context, nn, dataset as ds
import mindspore.dataset.vision.c_transforms as CV
import mindspore.dataset.transforms.c_transforms as C
from mindspore.train.callback import ModelCheckpoint, CheckpointConfig, Callback

import numpy as np
import os

context.set_context(mode=context.GRAPH_MODE, device_target="GPU")


class LeNet(nn.Cell):
    def __init__(self, num_classes=10):
        super(LeNet, self).__init__()

        self.conv1 = nn.Conv2d(1, 32, 5, pad_mode='valid')
        self.pool1 = nn.MaxPool2d(2, stride=2)
        self.conv2 = nn.Conv2d(32, 64, 5, pad_mode='valid')
        self.pool2 = nn.MaxPool2d(2, stride=2)

        self.act = nn.ReLU()
        self.flatten = nn.Flatten()

        self.fc1 = nn.Dense(1024, 512)
        self.dropout = nn.Dropout(0.5)
        self.out = nn.Dense(512, num_classes)

    def construct(self, x):
        x = self.pool1(self.act(self.conv1(x)))
        x = self.pool2(self.act(self.conv2(x)))

        x = self.flatten(x)

        x = self.fc1(x)
        x = self.dropout(x)
        return self.out(x)


def create_dataset(data_path, mode="train", batch_size=256, repeat_size=1, num_workers=4):
    # define dataset
    mnist_ds = ds.MnistDataset(os.path.join(data_path, mode))

    # using map to apply operations to a dataset
    if mode == "train":
        mnist_ds = mnist_ds.map(operations=CV.RandomResizedCrop((28, 28), scale=(0.8, 1.2)), input_columns="image", num_parallel_workers=num_parallel_workers)

    mnist_ds = mnist_ds.map(operations=C.TypeCast(mindspore.dtype.int32), input_columns="label", num_parallel_workers=num_workers)
    mnist_ds = mnist_ds.map(operations=CV.Rescale(1.0 / 255.0, 0), input_columns="image", num_parallel_workers=num_workers)
    mnist_ds = mnist_ds.map(operations=CV.HWC2CHW(), input_columns="image", num_parallel_workers=num_workers)

    # process the generated dataset
    if mode == "train":
        mnist_ds = mnist_ds.shuffle(buffer_size=10000)
    mnist_ds = mnist_ds.batch(batch_size, drop_remainder=True)
    mnist_ds = mnist_ds.repeat(repeat_size)

    return mnist_ds


class Solver(Callback):
    def __init__(self, epochs=100, batch_size=32, lr=0.01, data_dir="./data", num_workers=1):
        self.train_data = create_dataset(data_dir, "train", batch_size, 1, num_workers)
        self.val_data = create_dataset(data_dir, "test", batch_size, 1, num_workers)

        self.net = LeNet()
        self.loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction="mean")
        self.optimizer = nn.Momentum(self.net.trainable_params(), lr, 0.9, weight_decay=1e-4)
        self.model = mindspore.Model(self.net, self.loss, self.optimizer, metrics={"Accuracy": nn.Accuracy()})

        self.epochs = epochs
        self.output_dir = "output"
        os.makedirs(self.output_dir, exist_ok=True)

        self._loss_ls = []

    def train(self):
        config_ck = CheckpointConfig(save_checkpoint_steps=100, keep_checkpoint_max=16)
        ck_cb = ModelCheckpoint(prefix="lenet", directory=self.output_dir, config=config_ck)

        self.model.train(self.epochs, self.train_data, callbacks=[ck_cb, self], dataset_sink_mode=False)

    def epoch_begin(self, run_context):
        self._loss_ls.clear()

    def step_end(self, run_context):
        cb_params = run_context.original_args()
        loss = float(cb_params.net_outputs.asnumpy())

        self._loss_ls.append(loss)
        step = (cb_params.cur_step_num - 1) % cb_params.batch_num + 1
        print(f"[Epoch {cb_params.cur_epoch_num}/{cb_params.epoch_num}] "
              f"Step={step}/{cb_params.batch_num} "
              f"Loss={loss:.4f}/{np.mean(self._loss_ls):.4f}",
              end="\r", flush=True)

        if step == cb_params.batch_num:
            acc = self.model.eval(self.val_data, dataset_sink_mode=False)
            acc = float(acc["Accuracy"])
            print(f"\n[VAL] Accuracy={acc:.4f}\n")


if __name__ == '__main__':
    Solver().train()

赏

非常感谢,我会继续努力的！

支付宝

微信

• 本文作者： killf
• 本文链接： http://www.killf.info/机器学习/基于MindSpore的Mnist手写数字分类/
• 版权声明： 本博客所有文章除特别声明外，均采用 Apache License 2.0 许可协议。转载请注明出处！