本文章用的例子来自MindSpore官网教程，这里主要是分享一下个人理解和整合一下相关代码。

个人博客网站查看
环境配置：

• windows10
• MindSpore1.6.1 CPU版本
• python3.9.0

一、数据集和预训练模型下载

数据集下载

预训练模型下载

注：以上链接均来源MindSpore官网。

二、定义命令行参数并保存成config.yml

python">def parse_args():
    parser = argparse.ArgumentParser()
    
    # 工程文件名字
    parser.add_argument("--name", default="resnet50_classify", help="The name of project.")
    # 数据集名称
    parser.add_argument('--dataset', default='Canidae', help='dataset name')
    parser.add_argument("--epochs", default=200, type=int, metavar='N')
    parser.add_argument('--batch_size', default=16, type=int, metavar='N')

    # 预训练权重，参数填预训练文件的路径
    parser.add_argument('--pre_ckpt', default="./pre_ckpt/resnet50.ckpt")
    # 是否删除预训练模型的全连接层
    parser.add_argument("--delfc_flag", default=True)
    
    # 输入图片的channels，默认是RGB三通道图片
    parser.add_argument('--input_channels', default=3, type=int, help='input channels')
    # 类别个数
    parser.add_argument('--num_classes', default=20, type=int, help='number of classes')
    # 输入图像的尺寸
    parser.add_argument('--image_size', default=128, type=int, help='image size')

    # 优化器
    parser.add_argument('--optimizer', default='Adam')
    # 损失函数
    parser.add_argument('--loss', default='SoftmaxCrossEntropyWithLogits')
    
    parser.add_argument('--dataset_sink_mode', default=False)


    config = parser.parse_args()
    return config

使用argparse这个库可以自定义命令行参数，方便训练参数的更改。后面可以把这些参数存储为config.yml，作为配置文件使用，也方便自己或其他人了解这个网络的相关参数。

python">def main():
    # 解析命令行参数。config类型为字典，键是之前定义的命令行参数名字，值是对应其的内容。
    config = vars(parse_args())
    
    # 创建模型工程文件夹，用来保存训练文件
    if config["name"] is None:
        config['name'] = "test"
    os.makedirs(os.path.join("models", config["name"]), exist_ok=True)

    # 创建config文件并将参数写入
    config_path = 'models/%s/config.yml' % config['name']
    with open(config_path, 'w') as f:
        yaml.dump(config, f)

三、★数据集预处理和加载★

数据集分为训练集和验证集。

MindSpore提供了mindspore.dataset.ImageFolderDataset函数可以方便读取数据集。函数官方解释

不过要使用函数需要按照规定的方式存储数据。以本数据集为例，数据集可以以下面这种方式存放。

python">└─Canidae
  └─train
  │   └─dogs
  │   └─wolves
  └─val
      └─dogs
      └─wolves

该数据的存放方式是，在单个目录下创建要分类图片的对应文件夹，每一个文件夹对应一个类别。

例如，该函数会自动读取train目录下的所有文件夹，将同一个文件夹内的所有图片看作同一个类别。

在按文件夹名称顺序读入图片后会给每一类的图片赋予数字标签也就是label。第一个读入的文件夹图片被记为0，第二个读入的文件夹图片被记为1，依此类推直到所有文件夹读取完毕。

这种读取方式是函数默认的方式，但这种方式可能会让人无法准确判断哪一种类别对应着哪个label，在后面进行测试推理的时候可能会麻烦一点。因此我们可以使用该函数的class_indexing参数，用来直接指定文件夹与label的对应关系。

可以给这个参数传入字典类型。字典的键就是文件夹名称，值就是对应的数字。

例如：

python">dataset = ds.ImageFolderDataset(dataset_dir=image_folder_dataset_dir,
                                class_indexing={"dogs":0, "wolves":1})

这样就直接指定dogs是0，wolves是1。这就比较清晰label和类别的具体关系了。

因此我比较推荐直接将类别的名称作为文件夹的名字，可以通过函数读取文件夹的名字然后依次赋值生成相关字典并将内容保存到config文件中，这样之后测试就可以很清晰的知道类别和label关系。

python">def getClasses(data_path, config_path):
    """
    加载路径下所有文件夹的名字然后进行升序排序。
    将第一个文件夹记作0，第二个文件夹记作1，依此类推。
    然后将文件夹名称和对应的标记写入config文件，用于制作数据集label。
    Args:
        data_path (str): 数据集路径
        config_path (str): 配置文件路径
    Returns:
        _type_: _description_
    """
    res_dict = {"classes":{}}
    data_list = os.listdir(data_path)
    data_list.sort()
    for data in data_list:
        res_dict["classes"][data] = data_list.index(data)
    # 保存到config文件
    with open(config_path, 'a+') as f:
        yaml.dump(res_dict, f, Dumper=yaml.RoundTripDumper)
        f.close()
    return res_dict["classes"]

# 获取label，就是获取所有种类名称将其与数字对应
classesDict = getClasses(train_data_path, config_path)

在官方教程中给了create_dataset函数。这个函数是比较完整的数据集读取和处理。这个函数我就不多介绍了，能说的都写在下面的注释里了。

需要注意的是training参数。这个参数在True的时候是表示对训练集进行处理，False是对验证集进行处理。也就是只对训练集进行数据增强而验证集不用。

python">def create_dataset(data_path, image_size, classDict=None, batch_size=24, repeat_num=1, training=True):
    """定义数据集"""
    
    # 默认是按照文件夹名称排序（字母顺序），每一个类被赋予一个从0开始的唯一索引
    data_set = ds.ImageFolderDataset(data_path,
                                     num_parallel_workers=8,
                                     shuffle=True,
                                     class_indexing=classDict)

    # 对数据进行增强操作
    mean = [0.485 * 255, 0.456 * 255, 0.406 * 255]
    std = [0.229 * 255, 0.224 * 255, 0.225 * 255]
    if training:
        trans = [
            # 裁剪、Decode、Resize
            CV.RandomCropDecodeResize(image_size, scale=(0.8, 1.0), ratio=(0.75, 1.333)),
            # 水平翻转
            CV.RandomHorizontalFlip(prob=0.5),
            # 随机旋转
            CV.RandomRotation(90),
            # 归一化
            CV.Normalize(mean=mean, std=std),
            # h,w,c -> c,h,w
            CV.HWC2CHW()
        ]
    else:
        trans = [
            # 解码
            CV.Decode(),
            # 调整大小
            CV.Resize((image_size, image_size)),
            CV.Normalize(mean=mean, std=std),
            CV.HWC2CHW()
        ]
    type_cast_op = C.TypeCast(mstype.int32)
    # 实现数据的map映射、批量处理和数据重复的操作
    data_set = data_set.map(operations=trans, input_columns="image", num_parallel_workers=8)
    data_set = data_set.map(operations=type_cast_op, input_columns="label", num_parallel_workers=8)
    data_set = data_set.batch(batch_size, drop_remainder=True)
    data_set = data_set.repeat(repeat_num)

    return data_set

四、训练过程

主要的参数配置记录和数据集加载已经说完了，模型加载和训练方面MindSpore整合的很好，只需要几句话就可以调用，所以这部分就不多说了，在文章末尾会有工程的全部文件供下载。

五、测试过程

首先要加载分类标签，就是之前保存在config文件中的内容。

还需要将字典内容的键值调换顺序，方便后面根据预测结果直接输出类别名称。

python"># classes格式
classesDict = config["classes"]
# 将字典的键值调换顺序
class_name = dict(zip(classesDict.values(), classesDict.keys()))

★加载测试数据及预处理★

这一部分是很关键的。这一步的加载数据选择了与之前训练时不同的方式，是使用自定义的加载方式来读取数据集。

python">class testDataset:
    """
    自定义的读取图片的类。
    只用于test使用，不返回label。
    返回：
    第一个值为用于分类的图
    第二个值为该图的ID，即名字和后缀
    第三个值为用于显示的图
    """
    def __init__(self, data_path):
        self.data_path = data_path
        self.imgList = os.listdir(self.data_path)

    def __getitem__(self, index):
        imgID = self.imgList[index]
        img = cv2.imread(os.path.join(self.data_path, imgID))
        # 图像预处理
        img = cv2.resize(img, (256, 256))
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img = img.astype(np.float32)
        return img, imgID, img

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

# 读取测试数据并做预处理
testData = testDataset(test_data_path)
# image用于分类处理；imgID为图片名字和后缀；imageInit用于显示原图
dataset = ds.GeneratorDataset(testData, ["image", "imgID", "imageInit"])

# 图像预处理，主要是归一化和通道交换顺序
mean = [0.485 * 255, 0.456 * 255, 0.406 * 255]
std = [0.229 * 255, 0.224 * 255, 0.225 * 255]
dataset = dataset.map(operations=[CV.Normalize(mean=mean, std=std), CV.HWC2CHW()],
                      input_columns="image",
                      num_parallel_workers=8)

这一部分可能比较难理解。

首先自己定义了一个testDataset类去读取数据集并做简单的预处理，这个类的__getitem__函数返回三个值，分别是img，imgID，img。这三个值代表什么含义我在类的介绍里写了，可以去理解一下。

ds.GeneratorDataset这个函数可以创建一个dataset。第一个参数就填写之前写的testDataset类的实例结果，第二个参数填写“列名称”。这个“列名称”也就是column_names参数，我认为可以理解为通道名称。之前testDataset类的__getitem__函数可以返回三个值，这三个值可以通过ds.GeneratorDataset函数赋予每个值一个单独的名称，之后就可以通过这个“名称”调用对应的数据。

这部分可能难理解，我说的可能也不够详细和简洁，所以建议这部分单步调试，一点一点搞懂这个几个函数的用法比较好。

之后就是对图像进行归一化和通道数改变顺序的基本操作。

预测过程

python"># 用于存放预测结果
pred_list = []
# 存放显示图片
img_list = []
# 图片名字
imgID_list = []
# 扩充维度
expand = ExpandDims()
for data in dataset.create_dict_iterator():
    # 将图像添加到待显示图片列表中
    img_list.append(data["imageInit"].asnumpy())
    imgID_list.append(str(data["imgID"]))
    # 使用处理后的图像来预测
    image = data["image"].asnumpy()
    # 扩充维度
    img = expand(Tensor(image), 0)
    output = model.predict(img)
    pred = np.argmax(output.asnumpy(), axis=1)[0]
    pred_list.append(pred)

我这里的预测没有使用batch_size的格式，而是一张一张图片的预测，所以需要进行一个扩维的操作。

如果是batch_size格式的话需要在这之前对数据集进行batch处理。

模型最后的输出也就是pred是数字，这个数字与数据集的种类对应。具体的对应关系可以从class_name字典中获得。

python"># 可视化模型预测
plt.figure(figsize=(12, 5))
# 显示24张图
for i in range(len(img_list)):
    plt.subplot(3, 9, i+1)
    plt.title('{}'.format(class_name[pred_list[i]]))
    picture_show = img_list[i]/np.amax(img_list[i])
    picture_show = np.clip(picture_show, 0, 1)
    plt.imshow(picture_show)
    plt.axis('off')
plt.show()

可视化输出。可以显示预测结果和对应图片。

注意，这里显示的图片来自imageInit通道，因为image通道的图片都被归一化和通道顺序转换了所以没法直接用来显示。

可视化效果：

六、完整代码下载

Github下载

CSDN下载