MindSpore进行迁移学习

MindSpore进行迁移学习

学习自MindSpore官网案例，将ResNet50网络迁移学习至狼狗分类。此篇重点在模型的训练过程的代码上。

• 包含头文件

import mindspore as ms
import mindspore.dataset as ds
import mindspore.dataset.vision as vision # 用于数据增强
import matplotlib.pyplot as plt
import numpy as np
from typing import Type, Union, List, Optional
from mindspore import nn, train
from mindspore.common.initializer import Normal
import os
import time

• 加载数据集

batch_size = 18                             # 批量大小
image_size = 224                            # 训练图像空间大小
num_epochs = 10                             # 训练周期数
lr = 0.001                                  # 学习率
momentum = 0.9                              # momentum
workers = 4                                 # 并行线程个数

data_path_train = "./datasets/Canidae/train/"
data_path_val = "./datasets/Canidae/val/"

# 创建训练数据集
def create_dataset_canidae(dataset_path, usage):
    """ 数据加载 """
    data_set = ds.ImageFolderDataset(dataset_path,
                                     num_parallel_workers=workers,
                                     shuffle=True,)
    # 数据增强
    mean = [0.485 * 255, 0.456 * 255, 0.406 * 255]
    std = [0.229 * 255, 0.224 * 255, 0.225 * 233]
    scale = 32 # 缩放
    
    if usage == "train":
        trans = [
            vision.RandomCropDecodeResize(size=image_size, scale=(0.08,1.0), ratio=(0.75,1.333)),
            vision.RandomHorizontalFlip(prob=0.5),
            vision.Normalize(mean=mean, std=std),
            vision.HWC2CHW()
        ]
    else:
        trans = [
            vision.Decode(),
            vision.Resize(image_size + scale),
            vision.CenterCrop(image_size),
            vision.Normalize(mean=mean, std=std),
            vision.HWC2CHW()
        ]
    # 数据映射
    data_set = data_set.map(
        operations=trans,
        input_columns='image',
        num_parallel_workers=workers)
    data_set = data_set.batch(batch_size, drop_remainder=True)
    
    return data_set

dataset_train = create_dataset_canidae(data_path_train, "train")
step_size_train = dataset_train.get_dataset_size()

dataset_val = create_dataset_canidae(data_path_val, "val")
step_size_val = dataset_val.get_dataset_size()

这段代码的重点在于加载数据集函数中调用的ds.ImageFolderDataset函数，熟练此函数用法可用于自定义数据集。

• 要加载保存下来的权重文件，首先要在代码中实现一个可以作为权重数据载体的网络，在这里即是实现一个ResNet50网络。这里跳过此步骤，可在官网案例上查看。
• 这里假设已有ResNet()函数，用于构建模型。下面是加载预训练模型的方法：

def _resnet(block: Type[Union[ResidualBlockBase, ResidualBlock]],
            layers: List[int], num_classes: int = 1000, pretrained: bool, pretrianed_ckpt: str,
            input_channel: int):
    model = ResNet(block, layers, num_classes, input_channel)

    if pretrained:
        # 加载预训练模型
        param_dict = load_checkpoint(pretrianed_ckpt) # 这里是ckpt权重文件
        load_param_into_net(model, param_dict)

    return model

• 这里预训练的模型输出是1000个分类，我们只需要两个分类，因此要对网络进行更改。

network = resnet50(pretrained=True)

# 全连接层输入层的大小
in_channels = network.fc.in_channels
# 输出通道数大小为狼狗分类数2
head = nn.Dense(in_channels, 2)
# 重置全连接层
network.fc = head

# 平均池化层kernel size为7
avg_pool = nn.AvgPool2d(kernel_size=7)
# 重置平均池化层
network.avg_pool = avg_pool

注意，network.fc = head一行中，network.fc是网络中的全连接层，应按照网络中定义的名字(self.fc)进行相应的更改

• 接下来要定义优化器等，定义metrics的key值在训练时验证准确率会再用到一次。

opt = nn.Momentum(params=network.trainable_params(), learning_rate=lr, momentum=momentum)
loss_fn = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean')
metrics={"Accuracy": ms.nn.metrics.Accuracy()}

• 下一步，定义单步训练过程

def forward_fn(inputs, targets):

    logits = network(inputs)
    loss = loss_fn(logits, targets)

    return loss

grad_fn = ms.value_and_grad(forward_fn, None, opt.parameters)

def train_step(inputs, targets):

    loss, grads = grad_fn(inputs, targets)
    opt(grads)

    return loss

这里主要是两点，一个是grad_fn是依赖于forward_fn的，也包含了forward_fn，因此在train_step函数中只调用了grad_fn而没有forward_fn。另一个是这里的(inputs, targets)是对应的图像输入和标签。因此在forward_fn中，先把输入数据经过network后，将输出和标签进行损失函数的计算。
一般来说，此模板可套用。

• 接下来，创建迭代器，用于训练过程中加载数据集

# 创建迭代器
data_loader_train = dataset_train.create_tuple_iterator(num_epochs=num_epochs)
# 最佳模型保存路径
best_ckpt_dir = "./BestCheckpoint"
best_ckpt_path = "./BestCheckpoint/resnet50-best.ckpt"

• 准备工作完成，接下来就是开始训练模型部分了

# 开始循环训练
print("Start Training Loop ...")

best_acc = 0

for epoch in range(num_epochs):
    losses = []
    network.set_train()

    epoch_start = time.time()

    # 为每轮训练读入数据
    for i, (images, labels) in enumerate(data_loader_train):
        labels = labels.astype(ms.int32)
        loss = train_step(images, labels)
        losses.append(loss)

    # 每个epoch结束后，验证准确率

    acc = model.eval(dataset_val)['accuracy']

    epoch_end = time.time()
    epoch_seconds = (epoch_end - epoch_start) * 1000
    step_seconds = epoch_seconds/step_size_train

    print("-" * 20)
    print("Epoch: [%3d/%3d], Average Train Loss: [%5.3f], Accuracy: [%5.3f]" % (
        epoch+1, num_epochs, sum(losses)/len(losses), acc
    ))
    print("epoch time: %5.3f ms, per step time: %5.3f ms" % (
        epoch_seconds, step_seconds
    ))

    if acc > best_acc:
        best_acc = acc
        if not os.path.exists(best_ckpt_dir):
            os.mkdir(best_ckpt_dir)
        ms.save_checkpoint(network, best_ckpt_path)

print("=" * 80)
print(f"End of validation the best Accuracy is: {best_acc: 5.3f}, "
      f"save the best ckpt file in {best_ckpt_path}", flush=True)

• 这段代码重点在这几行：

	for i, (images, labels) in enumerate(data_loader_train):
        labels = labels.astype(ms.int32)
        loss = train_step(images, labels)
        losses.append(loss)

    # 每个epoch结束后，验证准确率

    acc = model.eval(dataset_val)['Accuracy']

注意acc这一行后面中括号的字符串要与metrics定义的字符串相同。