昇思25天学习打卡营第13天|基于MobileNetv2的垃圾分类

基于MobileNetv2的垃圾分类

本实验期望实现一个根据本地数据图像对垃圾物体检测，并将检测结果图片保存下来的应用，其中模型使用MobileNetv2。

背景知识 ——— MobileNetv2

MobileNet是google提出的专注于移动端、嵌入式和IoT设备的轻量级CNN网络。它能够在保持较高精准度的同时降低计算量和参数量。主要的创新点是倒残差结构和线性瓶颈。

1. 倒残差结构 Inverted Residuals
   传统的残差网络（ResNet）采用先增加通道数再减少通道数的结构，而MobileNetV2使用倒残差也就是先减少再增加。

这个结构由三部分组成

• 拓展卷积 Expansion Convolution
  使用1*1卷积将输入特征图的通道数增加到一个更高的维度，这里的维度由拓展因子决定，通常为6.
• 深度可分离卷积（Depthwise Separable Convolution）
  深度卷积：对每个输入通道都分别进行卷积，计算速度较快。
  逐点卷积：使用1*1卷积将通道数还原到原来的维度，结合所有通道的特征。
• 投影卷积 Projection Convolution
  使用1*1卷积将高维特征图投影回低纬度，生成输出特征图。

2. 线性瓶颈 Linear Bottlenecks
   在传统CNN中，非线性激活函数（例如ReLU）是较常用的增强模型非线性能力的激活函数。而在高维空间中使用ReLU可能会导致信息丢失。在MobileNet中，在特定的层中不使用ReLU激活函数，而是使用线性。线性瓶颈放在每个倒残差块的末尾，防止特征图的表示能力在降维过程中收到限制。

总体结构上，该网络由多个倒残差块和标准卷积层组成：

• 初始卷积层
  用于初步提取低级特征
• 倒残差块
  堆叠多个倒残差块，每个块都包含拓展、深度卷积和投影
• 全局平均池化层
  将特征图的空间维度降维倒1*1
• 全连接层
  输出分类结果

实验

数据处理

MobileNetV2默认使用的是ImageFolder的格式管理数据集。
每一类图片都是单独的一个文件夹。

from download import download
# 数据和权重文件 下载，依然还是从url中下载
url = "https://ascend-professional-construction-dataset.obs.cn-north-4.myhuaweicloud.com:443/MindStudio-pc/data_en.zip" 
# replace=True表示如果有同名文件会覆盖
path = download(url, "./", kind="zip", replace=True)

# 下载预训练权重文件
url = "https://ascend-professional-construction-dataset.obs.cn-north-4.myhuaweicloud.com:443/ComputerVision/mobilenetV2-200_1067.zip" 
path = download(url, "./", kind="zip", replace=True)

数据加载

os.environ['GLOG_v'] = '3' 
os.environ['GLOG_logtostderr'] = '0'
os.environ['GLOG_log_dir'] = '../../log' 
os.environ['GLOG_stderrthreshold'] = '2' 
set_context(mode=ms.GRAPH_MODE, device_target="CPU", device_id=0)

这段代码中设置了环境变量和执行上下文

• ‘GLOG_v’：设置日志级别，其中包括3(ERROR), 2(WARNING), 1(INFO), 0(DEBUG)，设为3仅记录错误消息。
• ‘GLOG_logtostderr’：指定日志的输出位置：包括0（将日志输出到文件），1（将日志输出到屏幕（标准错误输出））。
• GLOG_log_dir：制定日志目录
• GLOG_stderrthreshold：设置日志输出的标准错误阈值。即使日志输出到文件，当日志级别达到指定值时，也会输出到屏幕。这里设为二就是警告和错误信息都会输出到屏幕。

接下来设置了图模式的上下文，可以回顾下静态图加速。

然后定义后续会用到的参数。

# 垃圾分类数据集标签，以及用于标签映射的字典。
garbage_classes = {
    '干垃圾': ['贝壳', '打火机', '旧镜子', '扫把', '陶瓷碗', '牙刷', '一次性筷子', '脏污衣服'],
    '可回收物': ['报纸', '玻璃制品', '篮球', '塑料瓶', '硬纸板', '玻璃瓶', '金属制品', '帽子', '易拉罐', '纸张'],
    '湿垃圾': ['菜叶', '橙皮', '蛋壳', '香蕉皮'],
    '有害垃圾': ['电池', '药片胶囊', '荧光灯', '油漆桶']
}

class_cn = ['贝壳', '打火机', '旧镜子', '扫把', '陶瓷碗', '牙刷', '一次性筷子', '脏污衣服',
            '报纸', '玻璃制品', '篮球', '塑料瓶', '硬纸板', '玻璃瓶', '金属制品', '帽子', '易拉罐', '纸张',
            '菜叶', '橙皮', '蛋壳', '香蕉皮',
            '电池', '药片胶囊', '荧光灯', '油漆桶']
class_en = ['Seashell', 'Lighter','Old Mirror', 'Broom','Ceramic Bowl', 'Toothbrush','Disposable Chopsticks','Dirty Cloth',
            'Newspaper', 'Glassware', 'Basketball', 'Plastic Bottle', 'Cardboard','Glass Bottle', 'Metalware', 'Hats', 'Cans', 'Paper',
            'Vegetable Leaf','Orange Peel', 'Eggshell','Banana Peel',
            'Battery', 'Tablet capsules','Fluorescent lamp', 'Paint bucket']

index_en = {'Seashell': 0, 'Lighter': 1, 'Old Mirror': 2, 'Broom': 3, 'Ceramic Bowl': 4, 'Toothbrush': 5, 'Disposable Chopsticks': 6, 'Dirty Cloth': 7,
            'Newspaper': 8, 'Glassware': 9, 'Basketball': 10, 'Plastic Bottle': 11, 'Cardboard': 12, 'Glass Bottle': 13, 'Metalware': 14, 'Hats': 15, 'Cans': 16, 'Paper': 17,
            'Vegetable Leaf': 18, 'Orange Peel': 19, 'Eggshell': 20, 'Banana Peel': 21,
            'Battery': 22, 'Tablet capsules': 23, 'Fluorescent lamp': 24, 'Paint bucket': 25}

from easydict import EasyDict
# 训练超参
config = EasyDict({
    "num_classes": 26,
    "image_height": 224,
    "image_width": 224,
    "backbone_out_channels":1280,
    "batch_size": 16,
    "eval_batch_size": 8,
    "epochs": 10,
    "lr_max": 0.05,
    "momentum": 0.9,
    "weight_decay": 1e-4,
    "save_ckpt_epochs": 1,
    "dataset_path": "./data_en",
    "class_index": index_en,
    "pretrained_ckpt": "./mobilenetV2-200_1067.ckpt"  
})

EasyDict是一个可以将字典对象转换为类似面向对象方式访问的python库。

简单来介绍下这里的超参。

• num_classes：类别数量26，也就是需要分类的总类别数是26类。
• image_height/image_width：图片长度宽度
• backbone_out_channels：表示骨干网络（backbone）的输出的特征图的通道数
• batch_size：用于训练的批大小
• eval_batch_size：用于评估时的批大小
• epochs：对整个数据集的训练轮数
• lr_max：最大学习率
• momentum：动量优化器的动量参数，能帮助加速收敛和稳定训练过程。
• weight_decay：权重衰败系数，防止过拟合。
• save_ckpt_epochs：保存检查点的频率。设为1即为，一个周期保存一次模型检查点。
• dataset_path：数据集路径
• class_index：类别索引，将类标签映射到索引上
• pretrained_ckpt：预训练模型的检查点路径

数据预处理

使用mindSpore.dataset中的ImageFolderDataset读取数据集。下面的代码可以创建训练和评估数据集。

def create_dataset(dataset_path, config, training=True, buffer_size=1000):
    data_path = os.path.join(dataset_path, 'train' if training else 'test')
    ds = de.ImageFolderDataset(data_path, num_parallel_workers=4, class_indexing=config.class_index)
    resize_height = config.image_height
    resize_width = config.image_width
    
    normalize_op = C.Normalize(mean=[0.485*255, 0.456*255, 0.406*255], std=[0.229*255, 0.224*255, 0.225*255])
    change_swap_op = C.HWC2CHW()
    type_cast_op = C2.TypeCast(mstype.int32)
​
    if training:
        crop_decode_resize = C.RandomCropDecodeResize(resize_height, scale=(0.08, 1.0), ratio=(0.75, 1.333))
        horizontal_flip_op = C.RandomHorizontalFlip(prob=0.5)
        color_adjust = C.RandomColorAdjust(brightness=0.4, contrast=0.4, saturation=0.4)
    
        train_trans = [crop_decode_resize, horizontal_flip_op, color_adjust, normalize_op, change_swap_op]
        train_ds = ds.map(input_columns="image", operations=train_trans, num_parallel_workers=4)
        train_ds = train_ds.map(input_columns="label", operations=type_cast_op, num_parallel_workers=4)
        
        train_ds = train_ds.shuffle(buffer_size=buffer_size)
        ds = train_ds.batch(config.batch_size, drop_remainder=True)
    else:
        decode_op = C.Decode()
        resize_op = C.Resize((int(resize_width/0.875), int(resize_width/0.875)))
        center_crop = C.CenterCrop(resize_width)
        
        eval_trans = [decode_op, resize_op, center_crop, normalize_op, change_swap_op]
        eval_ds = ds.map(input_columns="image", operations=eval_trans, num_parallel_workers=4)
        eval_ds = eval_ds.map(input_columns="label", operations=type_cast_op, num_parallel_workers=4)
        ds = eval_ds.batch(config.eval_batch_size, drop_remainder=True)
​
    return ds

解释一下关键代码

ds = de.ImageFolderDataset(data_path, num_parallel_workers=4, class_indexing=config.class_index)
加载数据集，并行处理数为4，根据之前配置的超参设置类别索引。

后续都是一些图像的预处理操作，包括归一化，修改图像通道顺序，还有多标签进行类型转换。

若当前加载的是训练集，则还包括其他额外的数据增强。

• C.RandomCropDecodeResize 随机裁剪解码图像并调整大小
• RandomHorizontalFlip 随机对图像进行水平翻转
• RandomColorAdjust 随机调整图像的亮度、对比度和饱和度。
• train_trans = [crop_decode_resize, horizontal_flip_op, color_adjust, normalize_op, change_swap_op] 把所有操作组合成列表，后续使用map应用在image中。再对数据洗牌和批处理，drop_remainder 表示保留小于batch_size的批次。
  如果当前是训练集，则进行Decode解码，Resize调整图像大小，CenterCrop中心裁剪图片，最后进行分批。

处理后的图片大概长这样

搭建模型

使用mindspore.nn.Cell完成网络搭建。

__all__ = ['MobileNetV2', 'MobileNetV2Backbone', 'MobileNetV2Head', 'mobilenet_v2']

# 将v调整成可被divisor整除的最接近的值。
def _make_divisible(v, divisor, min_value=None):
    if min_value is None:
        min_value = divisor
    new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)
    if new_v < 0.9 * v:
        new_v += divisor
    return new_v

# 定义了平均池化层类
class GlobalAvgPooling(nn.Cell):
    def __init__(self):
        super(GlobalAvgPooling, self).__init__()

    def construct(self, x):
    # 前向传播逻辑是对x进行平均池化，对张量x在高度（维度2）和宽度（维度3）上取平均值
        x = P.mean(x, (2, 3))
        return x
        

定义卷积层、批量归一化、ReLU激活函数融合的块。详见注释。

class ConvBNReLU(nn.Cell):
   # 参数表示：输入通道数、输出通道数、卷积核大小、卷积步幅、通道分组数
    def __init__(self, in_planes, out_planes, kernel_size=3, stride=1, groups=1):
        super(ConvBNReLU, self).__init__()
        # 计算卷积核的填充
        padding = (kernel_size - 1) // 2
        in_channels = in_planes
        out_channels = out_planes
        if groups == 1:
        # 普通卷积层
            conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, pad_mode='pad', padding=padding)
        else:
        # 进行深度卷积，表示每个输入通道都有一个独立的卷积核。
            out_channels = in_planes
            conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, pad_mode='pad',
                             padding=padding, group=in_channels)
		# 构建 卷积层、批量归一化层、激活函数 的顺序容器
        layers = [conv, nn.BatchNorm2d(out_planes), nn.ReLU6()]
        self.features = nn.SequentialCell(layers)
	# 前向传播
    def construct(self, x):
        output = self.features(x)
        return output

定义倒残差块，详见注释。

class InvertedResidual(nn.Cell):
# 参数：输入通道数、输出通道数、卷积步幅、扩展比例（用于扩展输入通道数）
    def __init__(self, inp, oup, stride, expand_ratio):
        super(InvertedResidual, self).__init__()
        # 确认步幅值在1-2之间
        assert stride in [1, 2]
		# 计算隐藏通道数
        hidden_dim = int(round(inp * expand_ratio))
        # 当步幅为且输入输出通道数相等时，使用残差连接
        self.use_res_connect = stride == 1 and inp == oup

        layers = []
        # 若扩展比例不为1，添加ConvBNReLU层
        if expand_ratio != 1:
            layers.append(ConvBNReLU(inp, hidden_dim, kernel_size=1))
        # 添加ConvBNReLU层，Conv2d层（卷积层，将隐藏通道数压缩为输出通道数），BatchNorm2d（批量正则化层）
        layers.extend([
            ConvBNReLU(hidden_dim, hidden_dim,
                       stride=stride, groups=hidden_dim),
            nn.Conv2d(hidden_dim, oup, kernel_size=1,
                      stride=1, has_bias=False),
            nn.BatchNorm2d(oup),
        ])
        self.conv = nn.SequentialCell(layers)
        self.cast = P.Cast()
# 前向传播
    def construct(self, x):
        identity = x
        x = self.conv(x)
        # 若使用残差连接则返回identity（原始x）+ x（卷积后x）
        if self.use_res_connect:
            return P.add(identity, x)
        return x

定义MobileNetV2 的主干结构（backbone）

class MobileNetV2Backbone(nn.Cell):
    # 参数：通道数乘数（用于调整网络宽度）、倒置残差块的配置、将通道数四舍五入到最接近的倍数，输入通道数、输出通道数
    def __init__(self, width_mult=1., inverted_residual_setting=None, round_nearest=8,
                 input_channel=32, last_channel=1280):
        super(MobileNetV2Backbone, self).__init__()
        block = InvertedResidual
        # setting of inverted residual blocks
        self.cfgs = inverted_residual_setting
        # 如果没有倒置残差块的配置则使用基础配置
        if inverted_residual_setting is None:
            self.cfgs = [
                # t, c, n, s
                [1, 16, 1, 1],
                [6, 24, 2, 2],
                [6, 32, 3, 2],
                [6, 64, 4, 2],
                [6, 96, 3, 1],
                [6, 160, 3, 2],
                [6, 320, 1, 1],
            ]

        # 第一层ConvBNReLU
        input_channel = _make_divisible(input_channel * width_mult, round_nearest)
        self.out_channels = _make_divisible(last_channel * max(1.0, width_mult), round_nearest)
        # 将3通道转为输入的通道数
        features = [ConvBNReLU(3, input_channel, stride=2)]
        # 构建倒置残差块
        for t, c, n, s in self.cfgs:
            output_channel = _make_divisible(c * width_mult, round_nearest)
            for i in range(n):
                stride = s if i == 0 else 1
                features.append(block(input_channel, output_channel, stride, expand_ratio=t))
                input_channel = output_channel
        # 添加ConvBNReLU
        features.append(ConvBNReLU(input_channel, self.out_channels, kernel_size=1))
        self.features = nn.SequentialCell(features)
        self._initialize_weights()

    def construct(self, x):
        x = self.features(x)
        return x

    def _initialize_weights(self):
        # 初始化参数
        self.init_parameters_data()
        # 初始化卷积层和批量归一化层的参数
        # 卷积层权重使用正态分布
        # 批量归一化层gamma为1，beta为1
        for _, m in self.cells_and_names():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.set_data(Tensor(np.random.normal(0, np.sqrt(2. / n),
                                                          m.weight.data.shape).astype("float32")))
                if m.bias is not None:
                    m.bias.set_data(
                        Tensor(np.zeros(m.bias.data.shape, dtype="float32")))
            elif isinstance(m, nn.BatchNorm2d):
                m.gamma.set_data(
                    Tensor(np.ones(m.gamma.data.shape, dtype="float32")))
                m.beta.set_data(
                    Tensor(np.zeros(m.beta.data.shape, dtype="float32")))

    @property
    def get_features(self):
        return self.features

定义MobileNetv2的头部结构。

class MobileNetV2Head(nn.Cell):
 # 参数：输入通道数、输出类别数、是否dropout、激活函数类型（可选sigmoid和softmax）
    def __init__(self, input_channel=1280, num_classes=1000, has_dropout=False, activation="None"):
        super(MobileNetV2Head, self).__init__()
        # 如果dropout是false，头部包含的是全局平均池化层、全连接层
        # 不然就是全局平局池化层，dropout层、全连接层
        head = ([GlobalAvgPooling(), nn.Dense(input_channel, num_classes, has_bias=True)] if not has_dropout else
                [GlobalAvgPooling(), nn.Dropout(0.2), nn.Dense(input_channel, num_classes, has_bias=True)])
        self.head = nn.SequentialCell(head)
        self.need_activation = True
        if activation == "Sigmoid":
            self.activation = nn.Sigmoid()
        elif activation == "Softmax":
            self.activation = nn.Softmax()
        else:
            self.need_activation = False
        self._initialize_weights()

    def construct(self, x):
        x = self.head(x)
        if self.need_activation:
            x = self.activation(x)
        return x

    def _initialize_weights(self):
        self.init_parameters_data()
        for _, m in self.cells_and_names():
            if isinstance(m, nn.Dense):
            # 权重使用正态分布初始化，bias设为0
                m.weight.set_data(Tensor(np.random.normal(
                    0, 0.01, m.weight.data.shape).astype("float32")))
                if m.bias is not None:
                    m.bias.set_data(
                        Tensor(np.zeros(m.bias.data.shape, dtype="float32")))
    @property
    def get_head(self):
        return self.head

定义MobileNetV2的网络结构

class MobileNetV2(nn.Cell):
	# 参数：分类任务的类别数、通道数的乘数、是否dropout、是否有倒残差配置、通道数四舍五入的倍数、输入通道数、最终输出的通道数
    def __init__(self, num_classes=1000, width_mult=1., has_dropout=False, inverted_residual_setting=None, round_nearest=8, input_channel=32, last_channel=1280):
        super(MobileNetV2, self).__init__()
        # 定义主干结构
        self.backbone = MobileNetV2Backbone(width_mult=width_mult, \
            inverted_residual_setting=inverted_residual_setting, \
            round_nearest=round_nearest, input_channel=input_channel, last_channel=last_channel).get_features
        self.head = MobileNetV2Head(input_channel=self.backbone.out_channel, num_classes=num_classes, \
            has_dropout=has_dropout).get_head

    def construct(self, x):
    # 前向传播经过骨干结构和头部网络处理
        x = self.backbone(x)
        x = self.head(x)
        return x

# 将主干网络和头网络结合起来
class MobileNetV2Combine(nn.Cell):
    def __init__(self, backbone, head):
        super(MobileNetV2Combine, self).__init__(auto_prefix=False)
        self.backbone = backbone
        self.head = head

    def construct(self, x):
        x = self.backbone(x)
        x = self.head(x)
        return x
        
# 创建网络示例
def mobilenet_v2(backbone, head):
    return MobileNetV2Combine(backbone, head)

模型训练与测试

训练的进程到后期，模型会逐渐收敛，而模型参数的更新步幅也应相应逐渐降低，以减小模型抖动。由此训练时可以使用动态下降的学习率。这里使用了cosine decay下降策略。

def cosine_decay(total_steps, lr_init=0.0, lr_end=0.0, lr_max=0.1, warmup_steps=0):
	# 参数：总训练步数、初始学习率、结束学习率、学习率最大值、warmup步数
    lr_init, lr_end, lr_max = float(lr_init), float(lr_end), float(lr_max)
    # 计算衰退的步数
    decay_steps = total_steps - warmup_steps
    lr_all_steps = []
    # 计算线性的增量
    inc_per_step = (lr_max - lr_init) / warmup_steps if warmup_steps else 0
    for i in range(total_steps):
        if i < warmup_steps:
        # warmup期间线性增加
            lr = lr_init + inc_per_step * (i + 1)
        else:
        # 衰减期间，使用了余弦衰减公式
            cosine_decay = 0.5 * (1 + math.cos(math.pi * (i - warmup_steps) / decay_steps))
            lr = (lr_max - lr_end) * cosine_decay + lr_end
        lr_all_steps.append(lr)

    return lr_all_steps

保存模型

检查点checkpoint用于保存模型的参数，它可以保存模型结束时的参数，用以后续的推理。也可以保存中间的参数，防止模型训练异常退出后还要从最开始训练。此外，在finetuning微调场景下，还可以基于该模型对其他的任务进行模型的训练。

模型训练

本网络使用的损失函数是SoftmaxCrossEntropyWithLogits

from mindspore.amp import FixedLossScaleManager
import time
LOSS_SCALE = 1024

# 创建训练与评估数据集
train_dataset = create_dataset(dataset_path=config.dataset_path, config=config)
eval_dataset = create_dataset(dataset_path=config.dataset_path, config=config)
step_size = train_dataset.get_dataset_size()

# 实例化backbone
backbone = MobileNetV2Backbone() 
# 冻结backbone的参数
for param in backbone.get_parameters():
    param.requires_grad = False
# 加载预训练模型参数
load_checkpoint(config.pretrained_ckpt, backbone)
# 实例化头部
head = MobileNetV2Head(input_channel=backbone.out_channels, num_classes=config.num_classes)
# 构建完整的MobileNetV2模型
network = mobilenet_v2(backbone, head)

# 定义损失函数、学习率衰减、优化器
loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean')
loss_scale = FixedLossScaleManager(LOSS_SCALE, drop_overflow_update=False)
lrs = cosine_decay(config.epochs * step_size, lr_max=config.lr_max)
opt = nn.Momentum(network.trainable_params(), lrs, config.momentum, config.weight_decay, loss_scale=LOSS_SCALE)

# 定义用于训练的train_loop函数。
def train_loop(model, dataset, loss_fn, optimizer):
    # 定义正向计算函数
    def forward_fn(data, label):
        logits = model(data)
        loss = loss_fn(logits, label)
        return loss

    # 定义微分函数，使用mindspore.value_and_grad获得微分函数grad_fn,输出loss和梯度。
    # 由于是对模型参数求导,grad_position 配置为None，传入可训练参数。
    grad_fn = ms.value_and_grad(forward_fn, None, optimizer.parameters)

    # 定义 one-step training函数
    def train_step(data, label):
        loss, grads = grad_fn(data, label)
        optimizer(grads)
        return loss

    size = dataset.get_dataset_size()
    model.set_train()
    for batch, (data, label) in enumerate(dataset.create_tuple_iterator()):
        loss = train_step(data, label)

        if batch % 10 == 0:
            loss, current = loss.asnumpy(), batch
            print(f"loss: {loss:>7f}  [{current:>3d}/{size:>3d}]")

# 定义用于测试的test_loop函数。
def test_loop(model, dataset, loss_fn):
    num_batches = dataset.get_dataset_size()
    model.set_train(False)
    total, test_loss, correct = 0, 0, 0
    for data, label in dataset.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")

# 开始训练，每次训练都会评估
print("============== Starting Training ==============")
epoch_begin_time = time.time()
epochs = 2
for t in range(epochs):
    begin_time = time.time()
    print(f"Epoch {t+1}\n-------------------------------")
    train_loop(network, train_dataset, loss, opt)
    ms.save_checkpoint(network, "save_mobilenetV2_model.ckpt")
    end_time = time.time()
    times = end_time - begin_time
    print(f"per epoch time: {times}s")
    test_loop(network, eval_dataset, loss)
epoch_end_time = time.time()
times = epoch_end_time - epoch_begin_time
print(f"total time:  {times}s")
print("============== Training Success ==============")

模型推理

可以直接加载模型的checkpoint推理。要把数据传入网络，而不是给带了优化器和损失函数的网络。

CKPT="save_mobilenetV2_model.ckpt"
def image_process(image):
   	# 处理一个图片的张量
    mean=[0.485*255, 0.456*255, 0.406*255]
    std=[0.229*255, 0.224*255, 0.225*255]
    image = (np.array(image) - mean) / std
    image = image.transpose((2,0,1))
    img_tensor = Tensor(np.array([image], np.float32))
    return img_tensor

def infer_one(network, image_path):
# 推理一张图片
    image = Image.open(image_path).resize((config.image_height, config.image_width))
    logits = network(image_process(image))
    pred = np.argmax(logits.asnumpy(), axis=1)[0]
    print(image_path, class_en[pred])

def infer():
# 定义主干网络、头部网络，拼成一个完整网络
    backbone = MobileNetV2Backbone(last_channel=config.backbone_out_channels)
    head = MobileNetV2Head(input_channel=backbone.out_channels, num_classes=config.num_classes)
    network = mobilenet_v2(backbone, head)
    load_checkpoint(CKPT, network)
    for i in range(91, 100):
        infer_one(network, f'data_en/test/Cardboard/000{i}.jpg')
infer()

推理结果

导出模型文件

backbone = MobileNetV2Backbone(last_channel=config.backbone_out_channels)
head = MobileNetV2Head(input_channel=backbone.out_channels, num_classes=config.num_classes)
network = mobilenet_v2(backbone, head)
load_checkpoint(CKPT, network)

input = np.random.uniform(0.0, 1.0, size=[1, 3, 224, 224]).astype(np.float32)
# 三种导出的模型文件可选
# export(network, Tensor(input), file_name='mobilenetv2.air', file_format='AIR')
# export(network, Tensor(input), file_name='mobilenetv2.pb', file_format='GEIR')
export(network, Tensor(input), file_name='mobilenetv2.onnx', file_format='ONNX')

总结

本章完整地实现了MobileNetv2网络，并完成了模型的训练、保存、推理、导出。

打卡凭证