基于mindspore框架的数据集构建、模型搭建与训练

一、数据集构建

由于笔者的数据集是从Matlab中生成的结构体，并且已经命名好了“train_data”、“train_labels”、“test_data”和“test_labels”，并保存成了.mat文件，因此不同于像MNIST等现有的数据集可以直接下载，需要额外编写读取数据的代码。通过阅读mindspore官网，见下图：

可以构建一个类，并通过官方示例来自定义自己的数据集，整体流程是：1.先下载.mat文件并分别加载出训练测试数据和标签；2.在__getitem__方法中对数据类型进行转换（从numpy转到tensor，以及将维度转换为能够卷积的正确维度）；3.通过GeneratorDataset实现自定义数据集的加载。具体代码如下：

import numpy as np
from mindspore import Tensor
from scipy.io import loadmat
import mindspore as ms

class MyDataset:
    def __init__(self, data_path, use_type='train'):
        self.data_path = data_path
        self.data = []
        self.label = []

        # 加载 .mat 数据
        total_data = loadmat(data_path)
        if use_type == 'train':
            self.data = np.abs(total_data['train_data'])
            self.data = (self.data-self.data.min())/(self.data.max()-self.data.min())
            self.label = total_data['train_labels'].squeeze()
        elif use_type == 'test':
            self.data = np.abs(total_data['test_data'])
            self.data = (self.data-self.data.min())/(self.data.max()-self.data.min())
            self.label = total_data['test_labels'].squeeze()

    def __getitem__(self, idx):
        x = self.data[idx]
        y = self.label[idx]

        x = np.expand_dims(x.astype(np.float32), axis=0)
        y = np.int32(y)
        x = Tensor(x)
        y = Tensor(y)
        return x, y

    def __len__(self):
        return len(self.data)

def create_dataloader(data_path, batch_size, use_type='train'):
    ms.set_seed(1000)  # 设置随机种子以确保可重复性
    My_dataset = MyDataset(data_path=data_path, use_type=use_type)
    if use_type == 'train':
        My_dataloader = ms.dataset.GeneratorDataset(My_dataset, ["data", "labels"], shuffle=True)
    elif use_type == 'test':
        My_dataloader = ms.dataset.GeneratorDataset(My_dataset, ["data", "labels"], shuffle=False)  
    return My_dataloader.batch(batch_size)

二、网络搭建

由于笔者处理的是一维数据，因此搭建一个简单的1-D CNN，这里跟torch架构有所不同，比如torch的nn.Linear在mindspore下变成了nn.Dense；forward函数需要改为construct函数等，具体代码如下：

import math
import numpy as np
import mindspore as ms
import mindspore.nn as nn
from mindspore import Tensor, Parameter
from mindspore.common.initializer import XavierNormal, XavierUniform, Constant
import warnings

# CNN分类识别
class CNN_Net(nn.Cell):  # MindSpore中使用Cell替代PyTorch的Module
    def __init__(self, output_dim, seq_len):
        super(CNN_Net, self).__init__()
        self.cnn = nn.SequentialCell([  # 使用SequentialCell替代Sequential
            # Layer1
            nn.Conv1d(in_channels=1, out_channels=8, kernel_size=3, stride=1, has_bias=False,
                      weight_init=XavierNormal(), pad_mode='valid'),  # 添加权重初始化
            nn.BatchNorm1d(8),
            nn.Mish(),  # MindSpore同样支持Mish激活函数
            nn.MaxPool1d(kernel_size=2, stride=2),

            # Layer2
            nn.Conv1d(in_channels=8, out_channels=16, kernel_size=3, stride=1, has_bias=False,
                      weight_init=XavierNormal(), pad_mode='valid'),
            nn.BatchNorm1d(16),
            nn.Mish(),
            nn.MaxPool1d(kernel_size=2, stride=2),

            # Layer3
            nn.Conv1d(in_channels=16, out_channels=32, kernel_size=3, stride=1, has_bias=False,
                      weight_init=XavierNormal(), pad_mode='valid'),
            nn.BatchNorm1d(32),
            nn.Mish(),
            nn.MaxPool1d(kernel_size=2, stride=2)

        ])

        self.fc = nn.SequentialCell([
            nn.Dense(32*?, 16, weight_init=XavierNormal()),  # 使用Dense替代Linear，?为经过卷积和池化层后的维度，根据实际情况进行计算
            nn.Mish(),
            nn.Dense(16, output_dim, has_bias=False, weight_init=XavierNormal())
        ])

    def construct(self, x):  # MindSpore中使用construct替代forward
        x = self.cnn(x)
        x = x.view(x.shape[0], -1)  
        x = self.fc(x)
        return x
    

三、模型训练

训练时也跟torch框架下的代码有些出入，代码如下（笔者的代码是根据官方示例稍作修改后的，如果想要更加规范化可参照：模型训练 | MindSpore master 教程 | 昇思MindSpore社区）：

import numpy as np
import mindspore
import mindspore.nn as nn
from mindspore import Model
from mindspore import Tensor, context, ops
from mindspore.train.callback import LossMonitor
from mindspore.train.serialization import save_checkpoint
from sklearn.metrics import confusion_matrix
from tqdm import tqdm

from utils.data_utils import * 
from model.model import *

def forward_fn(data, label):
    logits = model(data)
    loss = loss_fn(logits, label)
    return loss, logits

if __name__ == '__main__':

    # 初始化参数
    seq_len = 32  # 数据长度
    batch_size = 128
    data_path = f'rcs_data/rcs_dataset_{seq_len}.mat'

    # 数据加载
    train_data = create_dataloader(data_path=data_path,batch_size=128,use_type='train')
    test_data = create_dataloader(data_path=data_path,batch_size=128,use_type='test')
    # for data in train_data:
    #     print(data)

    # 模型、损失、优化器
    learning_rate = 0.01
    model = CNN_Net(output_dim=2, seq_len=seq_len)
    loss_fn = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean')
    optimizer = nn.SGD(model.trainable_params(), learning_rate=learning_rate, weight_decay=1e-4)

    epochs = 10  # 训练轮数
    flag = 0
    grad_fn = mindspore.value_and_grad(forward_fn, None, optimizer.parameters, has_aux=True) # 使用value_and_grad通过函数变换获得梯度计算函数
    
    for epoch in range(epochs):
        print(f"<-----第{epoch+1}轮训练开始------>")
        # 训练过程
        size = train_data.get_dataset_size()
        model.set_train() # 设置为训练模式
        total, train_loss, correct = 0, 0, 0
        for batch, (data, label) in enumerate(train_data.create_tuple_iterator()):
            pred = model(data)
            total += len(data)
            correct += (pred.argmax(1) == label).asnumpy().sum()
            (loss, logits), grads = grad_fn(data, label)
            optimizer(grads)
            loss, current = loss.asnumpy(), batch
            train_loss += loss
        correct /= total
        print(f"Train: \n Accuracy: {(100*correct):>0.1f}%, Total loss: {train_loss:>7f} \n")

        # 测试过程
        num_batches = test_data.get_dataset_size()
        model.set_train(False)
        total, test_loss, correct = 0, 0, 0
        for data, label in test_data.create_tuple_iterator():
            pred = model(data)
            total += len(data)
            test_loss += loss_fn(pred, label).asnumpy()
            correct += (pred.argmax(1) == label).asnumpy().sum()
        test_loss /= num_batches
        correct /= total
        print(f"Test: \n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f} \n")
        #保存模型
        # file_name = f"model_weights/model_CNN_{seq_len}_{epoch+1}_{correct}.ckpt"
        # mindspore.save_checkpoint(model, file_name)
        if correct > flag:
            flag = correct
            print(f"模型准确率提升，保存模型，当前最高准确率为{(100*flag):>0.1f}")
            file_name = f"model_weights/model_CNN_{seq_len}_{epoch+1}_{(100*flag):>0.1f}.ckpt"
            mindspore.save_checkpoint(model, file_name)

PS：使用mindspore框架训练模型时还可以用更为高级的回调机制（Callback），感兴趣的朋友可以参考官方示例，个人感觉有些复杂而且对于比较简单的网络和训练过程没有必要使用回调机制~