当前案例不支持在GPU设备上静态图模式运行,其他模式运行皆支持。
1、ShuffleNet网络介绍
ShuffleNetV1是旷视科技提出的一种计算高效的CNN模型,和MobileNet, SqueezeNet等一样主要应用在移动端,所以模型的设计目标就是利用有限的计算资源来达到最好的模型精度。ShuffleNetV1的设计核心是引入了两种操作:Pointwise Group Convolution和Channel Shuffle,这在保持精度的同时大大降低了模型的计算量。因此,ShuffleNetV1和MobileNet类似,都是通过设计更高效的网络结构来实现模型的压缩和加速。
了解ShuffleNet更多详细内容,详见论文ShuffleNet。
如下图所示,ShuffleNet在保持不低的准确率的前提下,将参数量几乎降低到了最小,因此其运算速度较快,单位参数量对模型准确率的贡献非常高。
图片来源:Bianco S, Cadene R, Celona L, et al. Benchmark analysis of representative deep neural network architectures[J]. IEEE access, 2018, 6: 64270-64277.
1.1 模型架构
ShuffleNet最显著的特点在于对不同通道进行重排来解决Group Convolution带来的弊端。通过对ResNet的Bottleneck单元进行改进,在较小的计算量的情况下达到了较高的准确率。
Pointwise Group Convolution
Group Convolution(分组卷积)原理如下图所示,相比于普通的卷积操作,分组卷积的情况下,每一组的卷积核大小为in_channels/g*k*k,一共有g组,所有组共有(in_channels/g*k*k)*out_channels个参数,是正常卷积参数的1/g。分组卷积中,每个卷积核只处理输入特征图的一部分通道,其优点在于参数量会有所降低,但输出通道数仍等于卷积核的数量。
图片来源:Huang G, Liu S, Van der Maaten L, et al. Condensenet: An efficient densenet using learned group convolutions[C]//Proceedings of the IEEE conference on computer vision and pattern recognition. 2018: 2752-2761.
Depthwise Convolution(深度可分离卷积)将组数g分为和输入通道相等的in_channels
,然后对每一个in_channels
做卷积操作,每个卷积核只处理一个通道,记卷积核大小为1*k*k,则卷积核参数量为:in_channels*k*k,得到的feature maps通道数与输入通道数相等;
Pointwise Group Convolution(逐点分组卷积)在分组卷积的基础上,令每一组的卷积核大小为 1×11×1,卷积核参数量为(in_channels/g*1*1)*out_channels。
from mindspore import nn
import mindspore.ops as ops
from mindspore import Tensor
class GroupConv(nn.Cell):
#定义类的初始化方法(构造函数),接收以下参数:in_channels:输入通道数;out_channels:输出通道数;kernel_size:卷积核的大小
#stride:卷积的步长;pad_mode:填充模式,默认是 "pad";pad:填充大小,默认是0;groups:分组数,默认是1;has_bias:是否有偏置,默认是False
def __init__(self, in_channels, out_channels, kernel_size,
stride, pad_mode="pad", pad=0, groups=1, has_bias=False):
super(GroupConv, self).__init__()#调用父类 nn.Cell 的初始化方法
self.groups = groups#将分组数 groups 保存为类的实例变量
self.convs = nn.CellList()#创建一个 nn.CellList 实例 self.convs,用于存储多个卷积层。
for _ in range(groups):#遍历每一个分组
#每次遍历中,向 self.convs 添加一个新的卷积层,该卷积层具有以下特点:
#输入通道数和输出通道数分别是 in_channels // groups 和 out_channels // groups(即每组处理的通道数)。
# 其他参数如 kernel_size, stride, has_bias, padding, pad_mode 和 weight_init 使用传入的参数值。
self.convs.append(nn.Conv2d(in_channels // groups, out_channels // groups,
kernel_size=kernel_size, stride=stride, has_bias=has_bias,
padding=pad, pad_mode=pad_mode, group=1, weight_init='xavier_uniform'))
def construct(self, x):#定义 construct 方法,用于执行分组卷积操作。接收输入 x。
# 使用 ops.split 将输入 x 按照通道维度分割成多个部分,每个部分的大小为 len(x[0]) // self.groups。
features = ops.split(x, split_size_or_sections=int(len(x[0]) // self.groups), axis=1)
outputs = ()
for i in range(self.groups):
# 对每个分组,使用相应的卷积层处理分割后的输入部分,并将结果添加到 outputs 元组中。将输入部分转换为 float32 类型。
outputs = outputs + (self.convs[i](features[i].astype("float32")),)
#使用 ops.cat 将所有分组的卷积输出在通道维度上拼接在一起,得到最终输出 out。
out = ops.cat(outputs, axis=1)
return out
1.2 Channel Shuffle
Group Convolution的弊端在于不同组别的通道无法进行信息交流,堆积GConv层后一个问题是不同组之间的特征图是不通信的,这就好像分成了g个互不相干的道路,每一个人各走各的,这可能会降低网络的特征提取能力。这也是Xception,MobileNet等网络采用密集的1x1卷积(Dense Pointwise Convolution)的原因。
为了解决不同组别通道“近亲繁殖”的问题,ShuffleNet优化了大量密集的1x1卷积(在使用的情况下计算量占用率达到了惊人的93.4%),引入Channel Shuffle机制(通道重排)。这项操作直观上表现为将不同分组通道均匀分散重组,使网络在下一层能处理不同组别通道的信息。
如下图所示,对于g组,每组有n个通道的特征图,首先reshape成g行n列的矩阵,再将矩阵转置成n行g列,最后进行flatten操作,得到新的排列。这些操作都是可微分可导的且计算简单,在解决了信息交互的同时符合了ShuffleNet轻量级网络设计的轻量特征。
1.3 ShffleNet模块
如下图所示,ShuffleNet对ResNet中的Bottleneck结构进行由(a)到(b), (c)的更改:
-
将开始和最后的1×11×1卷积模块(降维、升维)改成Point Wise Group Convolution;
-
为了进行不同通道的信息交流,再降维之后进行Channel Shuffle;
-
降采样模块中,3×33×3 Depth Wise Convolution的步长设置为2,长宽降为原来的一般,因此shortcut中采用步长为2的3×33×3平均池化,并把相加改成拼接。
class ShuffleV1Block(nn.Cell):
# ShuffleV1Block 类的构造函数,接收以下参数:
# first_group: 是否在第一个卷积层使用分组卷积;mid_channels: 中间通道数;ksize: 卷积核大小;stride: 步长
def __init__(self, inp, oup, group, first_group, mid_channels, ksize, stride):
super(ShuffleV1Block, self).__init__()
self.stride = stride#步长 stride 保存为类的实例变量。
pad = ksize // 2
self.group = group
# 根据步长 stride 的值决定 outputs,如果步长为2,输出通道数为 oup - inp,否则为 oup。
if stride == 2:
outputs = oup - inp
else:
outputs = oup
self.relu = nn.ReLU()
# 定义第一个主分支 branch_main_1,包含以下层:
# 分组卷积层,输入通道为 inp,输出通道为 mid_channels,卷积核大小为1,步长为1,填充为0。是否使用分组卷积取决于 first_group。
# 批归一化层,处理 mid_channels。
# ReLU激活函数。
branch_main_1 = [
GroupConv(in_channels=inp, out_channels=mid_channels,
kernel_size=1, stride=1, pad_mode="pad", pad=0,
groups=1 if first_group else group),
nn.BatchNorm2d(mid_channels),
nn.ReLU(),
]
# 定义第二个主分支 branch_main_2,包含以下层:
# 卷积层,输入和输出通道均为 mid_channels,卷积核大小为 ksize,步长为 stride,填充为 pad,分组数为 mid_channels,使用 Xavier 权重初始化,不带偏置。
# 批归一化层,处理 mid_channels。
# 分组卷积层,输入通道为 mid_channels,输出通道为 outputs,卷积核大小为1,步长为1,填充为0。
# 批归一化层,处理 outputs。
branch_main_2 = [
nn.Conv2d(mid_channels, mid_channels, kernel_size=ksize, stride=stride,
pad_mode='pad', padding=pad, group=mid_channels,
weight_init='xavier_uniform', has_bias=False),
nn.BatchNorm2d(mid_channels),
GroupConv(in_channels=mid_channels, out_channels=outputs,
kernel_size=1, stride=1, pad_mode="pad", pad=0,
groups=group),
nn.BatchNorm2d(outputs),
]
# branch_main_1 和 branch_main_2 转换为 nn.SequentialCell,便于按顺序执行各层。
self.branch_main_1 = nn.SequentialCell(branch_main_1)
self.branch_main_2 = nn.SequentialCell(branch_main_2)
# 如果步长为2,定义一个平均池化层 branch_proj,卷积核大小为3,步长为2,填充模式为 'same'。
if stride == 2:
self.branch_proj = nn.AvgPool2d(kernel_size=3, stride=2, pad_mode='same')
# 定义 construct 方法,接收输入 old_x。将 old_x 分别赋值给 left 和 right 变量,初始化输出 out。
def construct(self, old_x):
left = old_x
right = old_x
out = old_x
# 将 right 通过第一个主分支 branch_main_1。
right = self.branch_main_1(right)
# 如果分组数大于1,对 right 进行通道洗牌操作。
if self.group > 1:
right = self.channel_shuffle(right)
# 将 right 通过第二个主分支 branch_main_2。
right = self.branch_main_2(right)
# 如果步长为1,将 left 和 right 相加,通过 ReLU 激活函数得到最终输出 out。
if self.stride == 1:
out = self.relu(left + right)
# 如果步长为2,将 left 通过平均池化层 branch_proj,然后将 left 和 right 在通道维度上拼接,通过 ReLU 激活函数得到最终输出 out。
elif self.stride == 2:
left = self.branch_proj(left)
out = ops.cat((left, right), 1)
out = self.relu(out)
return out
# 定义 channel_shuffle 方法,实现通道洗牌操作:
# 获取输入 x 的形状,分别为 batchsize、num_channels、height 和 width。
# 计算每个分组的通道数 group_channels。
# 将 x 重新调整形状为 (batchsize, group_channels, self.group, height, width)。
# 对 x 进行维度交换,使得分组维度排在前面。
# 将 x 重新调整回原始形状 (batchsize, num_channels, height, width)。
# 返回洗牌后的 x。
def channel_shuffle(self, x):
batchsize, num_channels, height, width = ops.shape(x)
group_channels = num_channels // self.group
x = ops.reshape(x, (batchsize, group_channels, self.group, height, width))
x = ops.transpose(x, (0, 2, 1, 3, 4))
x = ops.reshape(x, (batchsize, num_channels, height, width))
return x
2、构建ShuffleNet网络
ShuffleNet网络结构如下图所示,以输入图像224×224,组数3(g = 3)为例,首先通过数量24,卷积核大小为3×3,stride为2的卷积层,输出特征图大小为112×112,channel为24;然后通过stride为2的最大池化层,输出特征图大小为56×56,channel数不变;再堆叠3个ShuffleNet模块(Stage2, Stage3, Stage4),三个模块分别重复4次、8次、4次,其中每个模块开始先经过一次下采样模块(上图(c)),使特征图长宽减半,channel翻倍(Stage2的下采样模块除外,将channel数从24变为240);随后经过全局平均池化,输出大小为1×1×960,再经过全连接层和softmax,得到分类概率。
class ShuffleNetV1(nn.Cell):
# 定义类的初始化方法(构造函数),接收以下参数:
# n_class:分类的类别数,默认是1000。
# model_size:模型大小,默认是 '2.0x'。
# group:分组数,默认是3。
# 调用父类 nn.Cell 的初始化方法,并打印模型大小。
def __init__(self, n_class=1000, model_size='2.0x', group=3):
super(ShuffleNetV1, self).__init__()
print('model size is ', model_size)
# 定义每个阶段的重复次数,分别是4、8、4。
self.stage_repeats = [4, 8, 4]
# 将模型大小 model_size 保存为类的实例变量。
self.model_size = model_size
# 根据 group 和 model_size 的值,设置 self.stage_out_channels,这是每个阶段的输出通道数。每种 model_size 对应不同的通道配置。
if group == 3:
if model_size == '0.5x':
self.stage_out_channels = [-1, 12, 120, 240, 480]
elif model_size == '1.0x':
self.stage_out_channels = [-1, 24, 240, 480, 960]
elif model_size == '1.5x':
self.stage_out_channels = [-1, 24, 360, 720, 1440]
elif model_size == '2.0x':
self.stage_out_channels = [-1, 48, 480, 960, 1920]
else:
raise NotImplementedError
elif group == 8:
if model_size == '0.5x':
self.stage_out_channels = [-1, 16, 192, 384, 768]
elif model_size == '1.0x':
self.stage_out_channels = [-1, 24, 384, 768, 1536]
elif model_size == '1.5x':
self.stage_out_channels = [-1, 24, 576, 1152, 2304]
elif model_size == '2.0x':
self.stage_out_channels = [-1, 48, 768, 1536, 3072]
else:
raise NotImplementedError
# 设置第一个卷积层的输入通道数。
input_channel = self.stage_out_channels[1]
# 定义第一个卷积层,包括:
# 卷积层:输入通道数为3(RGB图像),输出通道数为 input_channel,卷积核大小为3,步长为2,填充为1,使用 Xavier 权重初始化,不带偏置。
# 批归一化层:处理 input_channel。
# ReLU激活函数。
self.first_conv = nn.SequentialCell(
nn.Conv2d(3, input_channel, 3, 2, 'pad', 1, weight_init='xavier_uniform', has_bias=False),
nn.BatchNorm2d(input_channel),
nn.ReLU(),
)
# 定义一个最大池化层,卷积核大小为3,步长为2,填充模式为 'same'。
self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, pad_mode='same')
# 定义特征提取部分。遍历每个阶段,根据 self.stage_repeats 添加多个 ShuffleV1Block:
# stride:如果是每个阶段的第一个块,步长为2,否则为1。
# first_group:如果是第一个阶段的第一个块,设置 first_group 为True。
# 将 ShuffleV1Block 添加到 features 列表中,并更新 input_channel 为 output_channel。
features = []
for idxstage in range(len(self.stage_repeats)):
numrepeat = self.stage_repeats[idxstage]
output_channel = self.stage_out_channels[idxstage + 2]
for i in range(numrepeat):
stride = 2 if i == 0 else 1
first_group = idxstage == 0 and i == 0
features.append(ShuffleV1Block(input_channel, output_channel,
group=group, first_group=first_group,
mid_channels=output_channel // 4, ksize=3, stride=stride))
input_channel = output_channel
# 将 features 列表转换为 nn.SequentialCell,便于按顺序执行各层。
self.features = nn.SequentialCell(features)
# 定义一个全局平均池化层,池化窗口大小为7。
self.globalpool = nn.AvgPool2d(7)
self.classifier = nn.Dense(self.stage_out_channels[-1], n_class)
# 定义 construct 方法,前向传播过程:
# 输入 x 通过第一个卷积层 self.first_conv。
# 通过最大池化层 self.maxpool。
# 通过特征提取部分 self.features。
# 通过全局平均池化层 self.globalpool。
# 将池化后的输出调整形状为二维,通道数为 self.stage_out_channels[-1]。
# 通过全连接层 self.classifier,得到最终的分类结果。
# 返回分类结果 x。
def construct(self, x):
x = self.first_conv(x)
x = self.maxpool(x)
x = self.features(x)
x = self.globalpool(x)
x = ops.reshape(x, (-1, self.stage_out_channels[-1]))
x = self.classifier(x)
return x
3、模型训练和评估
采用CIFAR-10数据集对ShuffleNet进行预训练。
3.1 训练集准备与加载
采用CIFAR-10数据集对ShuffleNet进行预训练。CIFAR-10共有60000张32*32的彩色图像,均匀地分为10个类别,其中50000张图片作为训练集,10000图片作为测试集。如下示例使用mindspore.dataset.Cifar10Dataset
接口下载并加载CIFAR-10的训练集。目前仅支持二进制版本(CIFAR-10 binary version)。
from download import download
url = "https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/notebook/datasets/cifar-10-binary.tar.gz"
download(url, "./dataset", kind="tar.gz", replace=True)
import mindspore as ms
from mindspore.dataset import Cifar10Dataset
from mindspore.dataset import vision, transforms
# 定义一个名为 get_dataset 的函数,接收三个参数:
# train_dataset_path: 训练数据集的路径
# batch_size: 每个批次的样本数量
# usage: 数据集的用途(例如 "train" 或 "test")
def get_dataset(train_dataset_path, batch_size, usage):
# 初始化一个空的 image_trans 列表。如果 usage 为 "train",则添加一系列数据增强和变换操作到 image_trans 列表中:
# 随机裁剪到 32x32,边距为 4 像素
# 随机水平翻转,概率为 0.5
# 调整图像大小到 224x224
# 像素值缩放到 [0, 1] 范围
# 使用均值和标准差进行归一化
# 将图像从 HWC 格式转换为 CHW 格式
image_trans = []
if usage == "train":
image_trans = [
vision.RandomCrop((32, 32), (4, 4, 4, 4)),
vision.RandomHorizontalFlip(prob=0.5),
vision.Resize((224, 224)),
vision.Rescale(1.0 / 255.0, 0.0),
vision.Normalize([0.4914, 0.4822, 0.4465], [0.2023, 0.1994, 0.2010]),
vision.HWC2CHW()
]
# 如果 usage 为 "test",则添加一系列数据变换操作到 image_trans 列表中(没有数据增强操作):
# 调整图像大小到 224x224
# 像素值缩放到 [0, 1] 范围
# 使用均值和标准差进行归一化
# 将图像从 HWC 格式转换为 CHW 格式
elif usage == "test":
image_trans = [
vision.Resize((224, 224)),
vision.Rescale(1.0 / 255.0, 0.0),
vision.Normalize([0.4914, 0.4822, 0.4465], [0.2023, 0.1994, 0.2010]),
vision.HWC2CHW()
]
# 定义标签变换操作,将标签转换为 ms.int32 类型。
label_trans = transforms.TypeCast(ms.int32)
# 加载CIFAR-10数据集,指定数据集路径和用途,并启用随机打乱。
dataset = Cifar10Dataset(train_dataset_path, usage=usage, shuffle=True)
# 对数据集应用变换操作:
dataset = dataset.map(image_trans, 'image')
dataset = dataset.map(label_trans, 'label')
# 将数据集分批次处理,每个批次包含 batch_size 个样本。如果最后一个批次样本数量不足,则丢弃该批次。
dataset = dataset.batch(batch_size, drop_remainder=True)
return dataset
# 调用 get_dataset 函数,加载训练数据集,路径为 ./dataset/cifar-10-batches-bin,每个批次包含 128 个样本。
dataset = get_dataset("./dataset/cifar-10-batches-bin", 128, "train")
# 获取每个epoch的批次数,并将其赋值给 batches_per_epoch 变量。
batches_per_epoch = dataset.get_dataset_size()
3.2 模型训练
本节用随机初始化的参数做预训练。首先调用ShuffleNetV1
定义网络,参数量选择"2.0x"
,并定义损失函数为交叉熵损失,学习率经过4轮的warmup
后采用余弦退火,优化器采用Momentum
。最后用train.model
中的Model
接口将模型、损失函数、优化器封装在model
中,并用model.train()
对网络进行训练。将ModelCheckpoint
、CheckpointConfig
、TimeMonitor
和LossMonitor
传入回调函数中,将会打印训练的轮数、损失和时间,并将ckpt文件保存在当前目录下。
import mindspore as ms
from mindspore.dataset import Cifar10Dataset
from mindspore.dataset import vision, transforms
# 定义一个名为 get_dataset 的函数,接收三个参数:
# train_dataset_path: 训练数据集的路径
# batch_size: 每个批次的样本数量
# usage: 数据集的用途(例如 "train" 或 "test")
def get_dataset(train_dataset_path, batch_size, usage):
# 初始化一个空的 image_trans 列表。如果 usage 为 "train",则添加一系列数据增强和变换操作到 image_trans 列表中:
# 随机裁剪到 32x32,边距为 4 像素
# 随机水平翻转,概率为 0.5
# 调整图像大小到 224x224
# 像素值缩放到 [0, 1] 范围
# 使用均值和标准差进行归一化
# 将图像从 HWC 格式转换为 CHW 格式
image_trans = []
if usage == "train":
image_trans = [
vision.RandomCrop((32, 32), (4, 4, 4, 4)),
vision.RandomHorizontalFlip(prob=0.5),
vision.Resize((224, 224)),
vision.Rescale(1.0 / 255.0, 0.0),
vision.Normalize([0.4914, 0.4822, 0.4465], [0.2023, 0.1994, 0.2010]),
vision.HWC2CHW()
]
# 如果 usage 为 "test",则添加一系列数据变换操作到 image_trans 列表中(没有数据增强操作):
# 调整图像大小到 224x224
# 像素值缩放到 [0, 1] 范围
# 使用均值和标准差进行归一化
# 将图像从 HWC 格式转换为 CHW 格式
elif usage == "test":
image_trans = [
vision.Resize((224, 224)),
vision.Rescale(1.0 / 255.0, 0.0),
vision.Normalize([0.4914, 0.4822, 0.4465], [0.2023, 0.1994, 0.2010]),
vision.HWC2CHW()
]
# 定义标签变换操作,将标签转换为 ms.int32 类型。
label_trans = transforms.TypeCast(ms.int32)
# 加载CIFAR-10数据集,指定数据集路径和用途,并启用随机打乱。
dataset = Cifar10Dataset(train_dataset_path, usage=usage, shuffle=True)
# 对数据集应用变换操作:
dataset = dataset.map(image_trans, 'image')
dataset = dataset.map(label_trans, 'label')
# 将数据集分批次处理,每个批次包含 batch_size 个样本。如果最后一个批次样本数量不足,则丢弃该批次。
dataset = dataset.batch(batch_size, drop_remainder=True)
return dataset
# 调用 get_dataset 函数,加载训练数据集,路径为 ./dataset/cifar-10-batches-bin,每个批次包含 128 个样本。
dataset = get_dataset("./dataset/cifar-10-batches-bin", 128, "train")
# 获取每个epoch的批次数,并将其赋值给 batches_per_epoch 变量。
batches_per_epoch = dataset.get_dataset_size()
训练好的模型保存在当前目录的shufflenetv1-5_390.ckpt
中,用作评估。
3.3 模型评估
在CIFAR-10的测试集上对模型进行评估。
设置好评估模型的路径后加载数据集,并设置Top 1, Top 5的评估标准,最后用model.eval()
接口对模型进行评估。
from mindspore import load_checkpoint, load_param_into_net
def test():
# 设置MindSpore的上下文为图模式,并使用Ascend设备
mindspore.set_context(mode=mindspore.GRAPH_MODE, device_target="Ascend")
# 加载CIFAR-10数据集,用于测试,batch size为128
dataset = get_dataset("./dataset/cifar-10-batches-bin", 128, "test")
# 初始化ShuffleNetV1模型,模型大小为"2.0x",分类数为10
net = ShuffleNetV1(model_size="2.0x", n_class=10)
# 从checkpoint文件中加载模型参数
param_dict = load_checkpoint("shufflenetv1-5_390.ckpt")
# 将加载的参数导入到网络中
load_param_into_net(net, param_dict)
# 设置模型为评估模式(非训练模式)
net.set_train(False)
# 定义交叉熵损失函数,使用标签平滑技术,损失减少方法为求平均值
loss = nn.CrossEntropyLoss(weight=None, reduction='mean', label_smoothing=0.1)
# 定义评估指标,包括损失、Top-1准确率和Top-5准确率
eval_metrics = {'Loss': nn.Loss(), 'Top_1_Acc': Top1CategoricalAccuracy(),
'Top_5_Acc': Top5CategoricalAccuracy()}
# 使用模型、损失函数和评估指标初始化Model对象
model = Model(net, loss_fn=loss, metrics=eval_metrics)
# 记录评估开始的时间
start_time = time.time()
# 进行模型评估,使用dataset_sink_mode=False表示不使用数据下沉模式
res = model.eval(dataset, dataset_sink_mode=False)
# 计算评估所用时间
use_time = time.time() - start_time
hour = str(int(use_time // 60 // 60))
minute = str(int(use_time // 60 % 60))
second = str(int(use_time % 60))
# 生成评估结果的日志信息
log = "result:" + str(res) + ", ckpt:'" + "./shufflenetv1-5_390.ckpt" \
+ "', time: " + hour + "h " + minute + "m " + second + "s"
print(log)
# 将日志信息写入到eval_log.txt文件中
filename = './eval_log.txt'#指定日志文件名为eval_log.txt。
with open(filename, 'a') as file_object:#以追加模式打开日志文件。
file_object.write(log + '\n')#将日志信息写入到文件中。
if __name__ == '__main__':
test()
3.4 模型预测
在CIFAR-10的测试集上对模型进行预测,并将预测结果可视化。
import mindspore
import matplotlib.pyplot as plt
import mindspore.dataset as ds
# 初始化ShuffleNetV1模型
net = ShuffleNetV1(model_size="2.0x", n_class=10)
# 创建一个空列表用于存储结果
show_lst = []
# 从checkpoint文件中加载模型参数
param_dict = load_checkpoint("shufflenetv1-5_390.ckpt")
load_param_into_net(net, param_dict)
# 使用加载的参数初始化模型
model = Model(net)
# 加载CIFAR-10数据集用于预测,设置不进行数据混洗
dataset_predict = ds.Cifar10Dataset(dataset_dir="./dataset/cifar-10-batches-bin", shuffle=False, usage="train")
# 加载CIFAR-10数据集用于显示,设置不进行数据混洗
dataset_show = ds.Cifar10Dataset(dataset_dir="./dataset/cifar-10-batches-bin", shuffle=False, usage="train")
dataset_show = dataset_show.batch(16)
# 获取一批图像用于显示
show_images_lst = next(dataset_show.create_dict_iterator())["image"].asnumpy()
# 定义图像转换操作
image_trans = [
vision.RandomCrop((32, 32), (4, 4, 4, 4)),
vision.RandomHorizontalFlip(prob=0.5),
vision.Resize((224, 224)),
vision.Rescale(1.0 / 255.0, 0.0),
vision.Normalize([0.4914, 0.4822, 0.4465], [0.2023, 0.1994, 0.2010]),
vision.HWC2CHW()
]
# 应用图像转换操作到预测数据集
dataset_predict = dataset_predict.map(image_trans, 'image')
dataset_predict = dataset_predict.batch(16)
# 定义类别字典,将类别标签转换为类别名称
class_dict = {0:"airplane", 1:"automobile", 2:"bird", 3:"cat", 4:"deer", 5:"dog", 6:"frog", 7:"horse", 8:"ship", 9:"truck"}
# 推理效果展示(上方为预测的结果,下方为推理效果图片)
# 创建图像显示窗口
plt.figure(figsize=(16, 5))
# 获取一批用于预测的数据
predict_data = next(dataset_predict.create_dict_iterator())
# 进行预测
output = model.predict(ms.Tensor(predict_data['image']))
# 获取预测结果的类别
pred = np.argmax(output.asnumpy(), axis=1)
# 显示预测结果和对应的图像
index = 0
for image in show_images_lst:
plt.subplot(2, 8, index+1)
plt.title('{}'.format(class_dict[pred[index]]))
index += 1
plt.imshow(image)
plt.axis("off")
plt.show()
打卡