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轻量级网络设计的轻量特征。
为了阅读方便,将Channel Shuffle的代码实现放在下方ShuffleNet模块的代码中。
ShuffleNet模块
如下图所示,ShuffleNet对ResNet中的Bottleneck结构进行由(a)到(b), (c)的更改:
-
将开始和最后的1×1卷积模块(降维、升维)改成Point Wise Group Convolution;
-
为了进行不同通道的信息交流,再降维之后进行Channel Shuffle;
-
降采样模块中,3×3 Depth Wise Convolution的步长设置为2,长宽降为原来的一半,因此shortcut中采用步长为2的3×3平均池化,并把相加改成拼接。
class ShuffleV1Block(nn.Cell):
def __init__(self, inp, oup, group, first_group, mid_channels, ksize, stride):
super(ShuffleV1Block, self).__init__()
self.stride = stride
pad = ksize // 2
self.group = group
if stride == 2:
outputs = oup - inp
else:
outputs = oup
self.relu = nn.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 = [
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),
]
self.branch_main_1 = nn.SequentialCell(branch_main_1)
self.branch_main_2 = nn.SequentialCell(branch_main_2)
if stride == 2:
self.branch_proj = nn.AvgPool2d(kernel_size=3, stride=2, pad_mode='same')
def construct(self, old_x):
left = old_x
right = old_x
out = old_x
right = self.branch_main_1(right)
if self.group > 1:
right = self.channel_shuffle(right)
right = self.branch_main_2(right)
if self.stride == 1:
out = self.relu(left + right)
elif self.stride == 2:
left = self.branch_proj(left)
out = ops.cat((left, right), 1)
out = self.relu(out)
return out
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 构建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):
def __init__(self, n_class=1000, model_size='2.0x', group=3):
super(ShuffleNetV1, self).__init__()
print('model size is ', model_size)
self.stage_repeats = [4, 8, 4]
self.model_size = 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]
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(),
)
self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, pad_mode='same')
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
self.features = nn.SequentialCell(features)
self.globalpool = nn.AvgPool2d(7)
self.classifier = nn.Dense(self.stage_out_channels[-1], n_class)
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 模型训练和评估
采用CIFAR-10数据集对ShuffleNet进行预训练。
训练集准备与加载
采用CIFAR-10数据集对ShuffleNet进行预训练。CIFAR-10共有60000张32*32的彩色图像,均匀地分为10个类别,其中50000张图片作为训练集,10000图片作为测试集。如下示例使用mindspore.dataset.Cifar10Dataset接口下载并加载CIFAR-10的训练集。目前仅支持二进制版本(CIFAR-10 binary version)。
模型训练
本节用随机初始化的参数做预训练。首先调用ShuffleNetV1定义网络,参数量选择"2.0x",并定义损失函数为交叉熵损失,学习率经过4轮的warmup后采用余弦退火,优化器采用Momentum。最后用train.model中的Model接口将模型、损失函数、优化器封装在model中,并用model.train()对网络进行训练。将ModelCheckpoint、CheckpointConfig、TimeMonitor和LossMonitor传入回调函数中,将会打印训练的轮数、损失和时间,并将ckpt文件保存在当前目录下。
训练好的模型保存在当前目录的shufflenetv1-250_391.ckpt中,用作评估。
模型评估
在CIFAR-10的测试集上对模型进行评估。
设置好评估模型的路径后加载数据集,并设置Top 1, Top 5的评估标准,最后用model.eval()接口对模型进行评估。
模型预测
在CIFAR-10的测试集上对模型进行预测,并将预测结果可视化。
from mindspore import nn import mindspore.ops as ops from mindspore import Tensor class GroupConv(nn.Cell): def __init__(self, in_channels, out_channels, kernel_size, stride, pad_mode="pad", pad=0, groups=1, has_bias=False): super(GroupConv, self).__init__() self.groups = groups self.convs = nn.CellList() for _ in range(groups): 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): features = ops.split(x, split_size_or_sections=int(len(x[0]) // self.groups), axis=1) outputs = () for i in range(self.groups): outputs = outputs + (self.convs[i](features[i].astype("float32")),) out = ops.cat(outputs, axis=1) return outChannel 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轻量级网络设计的轻量特征。
为了阅读方便,将Channel Shuffle的代码实现放在下方ShuffleNet模块的代码中。
ShuffleNet模块
如下图所示,ShuffleNet对ResNet中的Bottleneck结构进行由(a)到(b), (c)的更改:
将开始和最后的1×1卷积模块(降维、升维)改成Point Wise Group Convolution;
为了进行不同通道的信息交流,再降维之后进行Channel Shuffle;
降采样模块中,3×3 Depth Wise Convolution的步长设置为2,长宽降为原来的一半,因此shortcut中采用步长为2的3×3平均池化,并把相加改成拼接。
class ShuffleV1Block(nn.Cell): def __init__(self, inp, oup, group, first_group, mid_channels, ksize, stride): super(ShuffleV1Block, self).__init__() self.stride = stride pad = ksize // 2 self.group = group if stride == 2: outputs = oup - inp else: outputs = oup self.relu = nn.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 = [ 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), ] self.branch_main_1 = nn.SequentialCell(branch_main_1) self.branch_main_2 = nn.SequentialCell(branch_main_2) if stride == 2: self.branch_proj = nn.AvgPool2d(kernel_size=3, stride=2, pad_mode='same') def construct(self, old_x): left = old_x right = old_x out = old_x right = self.branch_main_1(right) if self.group > 1: right = self.channel_shuffle(right) right = self.branch_main_2(right) if self.stride == 1: out = self.relu(left + right) elif self.stride == 2: left = self.branch_proj(left) out = ops.cat((left, right), 1) out = self.relu(out) return out 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构建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): def __init__(self, n_class=1000, model_size='2.0x', group=3): super(ShuffleNetV1, self).__init__() print('model size is ', model_size) self.stage_repeats = [4, 8, 4] self.model_size = 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] 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(), ) self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, pad_mode='same') 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 self.features = nn.SequentialCell(features) self.globalpool = nn.AvgPool2d(7) self.classifier = nn.Dense(self.stage_out_channels[-1], n_class) 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模型训练和评估
采用CIFAR-10数据集对ShuffleNet进行预训练。
训练集准备与加载
采用CIFAR-10数据集对ShuffleNet进行预训练。CIFAR-10共有60000张32*32的彩色图像,均匀地分为10个类别,其中50000张图片作为训练集,10000图片作为测试集。如下示例使用
mindspore.dataset.Cifar10Dataset接口下载并加载CIFAR-10的训练集。目前仅支持二进制版本(CIFAR-10 binary version)。import mindspore as ms from mindspore.dataset import Cifar10Dataset from mindspore.dataset import vision, transforms def get_dataset(train_dataset_path, batch_size, usage): 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() ] 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() ] label_trans = transforms.TypeCast(ms.int32) dataset = Cifar10Dataset(train_dataset_path, usage=usage, shuffle=True) dataset = dataset.map(image_trans, 'image') dataset = dataset.map(label_trans, 'label') dataset = dataset.batch(batch_size, drop_remainder=True) return dataset dataset = get_dataset("./dataset/cifar-10-batches-bin", 128, "train") batches_per_epoch = dataset.get_dataset_size()模型训练
本节用随机初始化的参数做预训练。首先调用
ShuffleNetV1定义网络,参数量选择"2.0x",并定义损失函数为交叉熵损失,学习率经过4轮的warmup后采用余弦退火,优化器采用Momentum。最后用train.model中的Model接口将模型、损失函数、优化器封装在model中,并用model.train()对网络进行训练。将ModelCheckpoint、CheckpointConfig、TimeMonitor和LossMonitor传入回调函数中,将会打印训练的轮数、损失和时间,并将ckpt文件保存在当前目录下。import time import mindspore import numpy as np from mindspore import Tensor, nn from mindspore.train import ModelCheckpoint, CheckpointConfig, TimeMonitor, LossMonitor, Model, Top1CategoricalAccuracy, Top5CategoricalAccuracy def train(): mindspore.set_context(mode=mindspore.PYNATIVE_MODE, device_target="GPU") net = ShuffleNetV1(model_size="2.0x", n_class=10) loss = nn.CrossEntropyLoss(weight=None, reduction='mean', label_smoothing=0.1) min_lr = 0.0005 base_lr = 0.05 lr_scheduler = mindspore.nn.cosine_decay_lr(min_lr, base_lr, batches_per_epoch*250, batches_per_epoch, decay_epoch=250) lr = Tensor(lr_scheduler[-1]) optimizer = nn.Momentum(params=net.trainable_params(), learning_rate=lr, momentum=0.9, weight_decay=0.00004, loss_scale=1024) loss_scale_manager = ms.amp.FixedLossScaleManager(1024, drop_overflow_update=False) model = Model(net, loss_fn=loss, optimizer=optimizer, amp_level="O3", loss_scale_manager=loss_scale_manager) callback = [TimeMonitor(), LossMonitor()] save_ckpt_path = "./" config_ckpt = CheckpointConfig(save_checkpoint_steps=batches_per_epoch, keep_checkpoint_max=5) ckpt_callback = ModelCheckpoint("shufflenetv1", directory=save_ckpt_path, config=config_ckpt) callback += [ckpt_callback] print("============== Starting Training ==============") start_time = time.time() model.train(250, dataset, callbacks=callback) 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)) print("total time:" + hour + "h " + minute + "m " + second + "s") print("============== Train Success ==============") if __name__ == '__main__': train()训练好的模型保存在当前目录的
shufflenetv1-250_391.ckpt中,用作评估。模型评估
在CIFAR-10的测试集上对模型进行评估。
设置好评估模型的路径后加载数据集,并设置Top 1, Top 5的评估标准,最后用
model.eval()接口对模型进行评估。from mindspore import load_checkpoint, load_param_into_net def test(): mindspore.set_context(mode=mindspore.GRAPH_MODE, device_target="GPU") dataset = get_dataset("./dataset/cifar-10-batches-bin", 128, "test") net = ShuffleNetV1(model_size="2.0x", n_class=10) param_dict = load_checkpoint("shufflenetv1-250_391.ckpt") load_param_into_net(net, param_dict) net.set_train(False) loss = nn.CrossEntropyLoss(weight=None, reduction='mean', label_smoothing=0.1) eval_metrics = {'Loss': nn.Loss(), 'Top_1_Acc': Top1CategoricalAccuracy(), 'Top_5_Acc': Top5CategoricalAccuracy()} model = Model(net, loss_fn=loss, metrics=eval_metrics) start_time = time.time() 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-250_391.ckpt" \ + "', time: " + hour + "h " + minute + "m " + second + "s" print(log) filename = './eval_log.txt' with open(filename, 'a') as file_object: file_object.write(log + '\n') if __name__ == '__main__': test()model size is 2.0x result:{'Loss': 1.0217913215673422, 'Top_1_Acc': 0.8152, 'Top_5_Acc': 0.975}, ckpt:'./shufflenetv1-250_391.ckpt', time: 0h 0m 21s模型预测
在CIFAR-10的测试集上对模型进行预测,并将预测结果可视化。
import mindspore import matplotlib.pyplot as plt import mindspore.dataset as ds net = ShuffleNetV1(model_size="2.0x", n_class=10) show_lst = [] param_dict = load_checkpoint("shufflenetv1-250_391.ckpt") load_param_into_net(net, param_dict) model = Model(net) dataset_predict = ds.Cifar10Dataset(dataset_dir="./dataset/cifar-10-batches-bin", shuffle=False, usage="train") 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()