mindspore模型训练和模型导出为onnx

mindspore是华为深度学习框架，网址为：https://www.mindspore.cn/

本代码主要参考快速入门的代码，加了模型导出为onnx

mindspore在模型搭建上基本上的语法和pytorch差不多

只是分为了网络和模型，模型主要拿来训练和预测，而网络就是单纯的网络，网络可以拿来导出模型文件，但是预测只能使用模型

训练代码如下：

# -*- coding: utf-8 -*-
import os
import mindspore.nn as nn
from mindspore.common.initializer import Normal
from mindspore import Tensor, Model,export,load_checkpoint
from mindspore.nn.metrics import Accuracy
from mindspore.train.callback import ModelCheckpoint, CheckpointConfig, LossMonitor
from mindspore.nn.loss import SoftmaxCrossEntropyWithLogits
import mindspore.dataset.vision.c_transforms as CV
import mindspore.dataset.transforms.c_transforms as C
from mindspore.dataset.vision import Inter
from mindspore import dtype as mstype
import numpy as np
import mindspore.dataset as ds
from mindspore.train.callback import Callback


# https://obs.dualstack.cn-north-4.myhuaweicloud.com/mindspore-website/notebook/datasets/MNIST_Data.zip
train_data_path = "./datasets/MNIST_Data/train"
test_data_path = "./datasets/MNIST_Data/test"
mnist_path = "./datasets/MNIST_Data"
model_path = "./models/ckpt/"

#定义数据集
def create_dataset(data_path, batch_size=128, repeat_size=1,
                   num_parallel_workers=1):
    """ 
    create dataset for train or test
    
    Args:
        data_path (str): Data path
        batch_size (int): The number of data records in each group
        repeat_size (int): The number of replicated data records
        num_parallel_workers (int): The number of parallel workers
    """
    # define dataset
    mnist_ds = ds.MnistDataset(data_path)

    # define some parameters needed for data enhancement and rough justification
    resize_height, resize_width = 32, 32
    rescale = 1.0 / 255.0
    shift = 0.0
    rescale_nml = 1 / 0.3081
    shift_nml = -1 * 0.1307 / 0.3081

    # according to the parameters, generate the corresponding data enhancement method
    resize_op = CV.Resize((resize_height, resize_width), interpolation=Inter.LINEAR)
    rescale_nml_op = CV.Rescale(rescale_nml, shift_nml)
    rescale_op = CV.Rescale(rescale, shift)
    hwc2chw_op = CV.HWC2CHW()
    type_cast_op = C.TypeCast(mstype.int32)

    # using map to apply operations to a dataset
    mnist_ds = mnist_ds.map(operations=type_cast_op, input_columns="label", num_parallel_workers=num_parallel_workers)
    mnist_ds = mnist_ds.map(operations=resize_op, input_columns="image", num_parallel_workers=num_parallel_workers)
    mnist_ds = mnist_ds.map(operations=rescale_op, input_columns="image", num_parallel_workers=num_parallel_workers)
    mnist_ds = mnist_ds.map(operations=rescale_nml_op, input_columns="image", num_parallel_workers=num_parallel_workers)
    mnist_ds = mnist_ds.map(operations=hwc2chw_op, input_columns="image", num_parallel_workers=num_parallel_workers)
    
    # process the generated dataset
    buffer_size = 10000
    mnist_ds = mnist_ds.shuffle(buffer_size=buffer_size)
    mnist_ds = mnist_ds.batch(batch_size, drop_remainder=True)
    mnist_ds = mnist_ds.repeat(repeat_size)

    return mnist_ds

# custom callback function
class StepLossAccInfo(Callback):
    def __init__(self, model, eval_dataset, step_loss, steps_eval):
        self.model = model
        self.eval_dataset = eval_dataset
        self.step_loss = step_loss
        self.steps_eval = steps_eval
        
    def step_end(self, run_context):
        cb_params = run_context.original_args()
        cur_epoch = cb_params.cur_epoch_num
        cur_step = (cur_epoch-1)*1875 + cb_params.cur_step_num
        self.step_loss["loss_value"].append(str(cb_params.net_outputs))
        self.step_loss["step"].append(str(cur_step))
        if cur_step % 125 == 0:
            acc = self.model.eval(self.eval_dataset, dataset_sink_mode=False)
            self.steps_eval["step"].append(cur_step)
            self.steps_eval["acc"].append(acc["Accuracy"])

#定义网络            
class  mnist(nn.Cell):
     def __init__(self, num_class=10):
        super(mnist, self).__init__()
        self.conv1 = nn.Conv2d(1, 8, 5, pad_mode='valid')
        self.conv2 = nn.Conv2d(8, 12, 5, pad_mode='valid')
        self.fc1 = nn.Dense(300 , 120, weight_init=Normal(0.02))
        self.fc2 = nn.Dense(120, 60, weight_init=Normal(0.02))
        self.fc3 = nn.Dense(60, num_class, weight_init=Normal(0.02))
        self.relu = nn.ReLU()
        self.max_pool2d = nn.MaxPool2d(kernel_size=2, stride=2)
        self.flatten = nn.Flatten()
        
     def construct(self, x):
        x = self.max_pool2d(self.relu(self.conv1(x)))
        x = self.max_pool2d(self.relu(self.conv2(x)))        
        x = self.flatten(x)
        x = self.relu(self.fc1(x))
        x = self.relu(self.fc2(x))
        x = self.fc3(x) 
        return x    
                
network = mnist()
net_opt = nn.Momentum(network.trainable_params(), 0.01, 0.9)
net_loss = SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean')

#定义模型
model = Model(network, net_loss, net_opt, metrics={"Accuracy": Accuracy()} )
# save the network model and parameters for subsequence fine-tuning
config_ck = CheckpointConfig(save_checkpoint_steps=375, keep_checkpoint_max=16)
# group layers into an object with training and evaluation features
ckpoint_cb = ModelCheckpoint(prefix="mnist", directory=model_path, config=config_ck)
eval_dataset = create_dataset("./datasets/MNIST_Data/test")

step_loss = {"step": [], "loss_value": []}
steps_eval = {"step": [], "acc": []}
# collect the steps,loss and accuracy information
step_loss_acc_info = StepLossAccInfo(model , eval_dataset, step_loss, steps_eval)

repeat_size = 1
ds_train = create_dataset(os.path.join(mnist_path, "train"), 32, repeat_size)
model.train(1, ds_train, callbacks=[ckpoint_cb, LossMonitor(125), step_loss_acc_info], dataset_sink_mode=False)

#测试
ds_test = create_dataset(test_data_path).create_dict_iterator()
data = next(ds_test)
images = data["image"].asnumpy()
labels = data["label"].asnumpy()
print(model.predict(Tensor(data['image'])))
print(images.shape)
#导出格式为onnx文件
load_checkpoint("models\\ckpt\\mnist-1_1875.ckpt", net=network)
print("load****")
input =np.ones([1, 1, 32, 32]).astype(np.float32)
export(network, Tensor(input), file_name='minist', file_format='ONNX')

注意，数据集需要自己提前下载：https://obs.dualstack.cn-north-4.myhuaweicloud.com/mindspore-website/notebook/datasets/MNIST_Data.zip

运行后就可以导出模型

运行结果：

可以打开onnx文件看看网络结构：

 

验证一下数据结果:

onnx:

import onnxruntime
import numpy as np

x=np.ones([1, 1, 32, 32]).astype(np.float32)
session = onnxruntime.InferenceSession("minist.onnx")
inputs = {session.get_inputs()[0].name: x}
outs = session.run(None, inputs)
print('onnx result is:',outs)

运行结果：

 mindspore：

import mindspore.nn as nn
from mindspore.common.initializer import Normal
from mindspore import load_checkpoint, load_param_into_net
from mindspore import Tensor, Model
import numpy as np

class  mnist(nn.Cell):
     def __init__(self, num_class=10):
        super(mnist, self).__init__()
        self.conv1 = nn.Conv2d(1, 8, 5, pad_mode='valid')
        self.conv2 = nn.Conv2d(8, 12, 5, pad_mode='valid')
        self.fc1 = nn.Dense(300 , 120, weight_init=Normal(0.02))
        self.fc2 = nn.Dense(120, 60, weight_init=Normal(0.02))
        self.fc3 = nn.Dense(60, num_class, weight_init=Normal(0.02))
        self.relu = nn.ReLU()
        self.max_pool2d = nn.MaxPool2d(kernel_size=2, stride=2)
        self.flatten = nn.Flatten()
        
     def construct(self, x):
        x = self.max_pool2d(self.relu(self.conv1(x)))
        x = self.max_pool2d(self.relu(self.conv2(x)))        
        x = self.flatten(x)
        x = self.relu(self.fc1(x))
        x = self.relu(self.fc2(x))
        x = self.fc3(x) 
        return x
    
network = mnist()
ckpt_path = "models\\ckpt\\mnist-1_1875.ckpt"
trained_ckpt = load_checkpoint(ckpt_path)
load_param_into_net(network, trained_ckpt)
input=np.ones([1, 1, 32, 32]).astype(np.float32)
model = Model(network, metrics={'acc'}, eval_network=network)
print(model.predict(Tensor(input)))

运行结果：

运行结果一致，说明导出文件正确。