简单的 ConvMixer 会给 CV 带来新范式吗？
共 10077字，需浏览 21分钟
·
2021-11-27 04:42
1Patch 就够了？
近年来，卷积神经网络一直是视觉任务的主要架构，但最近的实验表明，基于 Transformer 的模型，尤其是 Vision Transformer (ViT)，在某些情况下可能会超过 CNN 的性能。
然而，由于 Transformer 中自注意力层的计算复杂度是关于 patch 数量的二次方，对于大图的计算量可观，因此后续一大波工作从这个角度作了一些改进，关于这点这里就不谈了。
ViT 的成功同时也带来了一个问题，那就是它的性能主要是由于 Transformer 架构（自注意力机制）的强大引起的，或者还是由于使用了 patch 作为输入表示主导的呢？而本篇的主角，也是 ICLR 2022 还在审稿中的新作 ConvMixer，就为后者提供了一些证据。
具体而言，ConvMixer 是一个极其简单的模型，在架构精神上与 ViT 以及更基本的 MLP-Mixer 相似，它也是直接将 patch 作为输入进行操作，但是它分离了空间和通道两个维度上的混合，并在整个网络中保持相同的通道数和分辨率。
所谓 ConvMixer，就是仅使用卷积来实现混合步骤。尽管它很简单，但作者表明 ConvMixer 的性能甚至优于 ViT、MLP-Mixer 及其类似的变体，此外还优于 ResNet 等经典视觉模型。
Patches Are All You Need从这个题目以及模型名字 ConvMixer 似乎可以感觉到，ViT 中的自注意力并不是必须的，只要使用合适大小的 patch，再加上通道内以及通道间的分离卷积混合，照样能抓取像素之间的远程关联，实现很好的数据表示。
该论文尚在评审中，有兴趣的可以前往参观 https://openreview.net/forum?id=TVHS5Y4dNvM。有些人认为该论文似乎并没有提供非常大的洞见以及理论，更偏向于从实验中发现了好用的结构设计并给以大家一定的启示。
2动机
ConvMixer 架构基于 patch 以及混合的基本思想。具体来说，
用 depthwise 卷积来混合通道内的值， 用 pointwise 卷积来混合通道间的值。
以往工作的一个关键思想是 MLP 和自注意力可以混合较远的空间信息，即它们可以具有任意大的感受野。因此，该研究通过使用较大的卷积核来实现混合远程关联。
3主要参数
ConvMixer 的实例化取决于四个参数：
1、patch 大小 ； 2、patch 的嵌入维度 ； 3、深度 ，即 ConvMixer 层的重复次数； 4、卷积层的 kernel 大小 。
可以根据它们的嵌入维度和深度命名具体的 ConvMixers，如 ConvMixer-h/d。
4ConvMixer 核心内容
参看上图，模型 ConvMixer 是由一个 patch 嵌入层以及一个简单的全卷积块的循环所构成。patch 大小为
ConvMixer 块本身由 depthwise 卷积（即，组数等于通道数 h 的分组卷积）和 pointwise（即 kernel 大小为 1×1）卷积组成。
每个卷积之后是一个激活和 BatchNorm：
在多次应用这个块之后，执行全局池化以获得大小为
5实 验
from tensorflow.keras import layers
from tensorflow import keras
import matplotlib.pyplot as plt
import tensorflow_addons as tfa
import tensorflow as tf
import numpy as np
+超参数
learning_rate = 0.001
weight_decay = 0.0001
batch_size = 128
num_epochs = 10
为了快速看到结果,模型仅仅训练 10 个 epoch，但后面可以看到，结果还是可以的。
+加载 CIFAR-10 数据集
(x_train, y_train), (x_test, y_test) = keras.datasets.cifar10.load_data()
val_split = 0.1
val_indices = int(len(x_train) * val_split)
new_x_train, new_y_train = x_train[val_indices:], y_train[val_indices:]
x_val, y_val = x_train[:val_indices], y_train[:val_indices]
print(f"Training data samples: {len(new_x_train)}")
print(f"Validation data samples: {len(x_val)}")
print(f"Test data samples: {len(x_test)}")
Training data samples: 45000
Validation data samples: 5000
Test data samples: 10000
+数据增强
image_size = 32
auto = tf.data.AUTOTUNE
data_augmentation = keras.Sequential(
[layers.RandomCrop(image_size, image_size), layers.RandomFlip('horizontal'),],
name='data_augmentation',
)
def make_datasets(images, labels, is_train=False):
dataset = tf.data.Dataset.from_tensor_slices((images, labels))
if is_train:
dataset = dataset.shuffle(batch_size * 10)
dataset = dataset.batch(batch_size)
if is_train:
dataset = dataset.map(
lambda x, y: (data_augmentation(x), y), num_parallel_calls=auto
)
return dataset.prefetch(auto)
train_dataset = make_datasets(new_x_train.astype(np.float32), new_y_train, is_train=True)
val_dataset = make_datasets(x_val.astype(np.float32), y_val)
test_dataset = make_datasets(x_test.astype(np.float32), y_test)
6网络实现
再次查看此图，我们直接根据这个流程图中的几个步骤来撸代码。
+1、计算 patch 的嵌入
块大小为
def conv_stem(x, filters: int, patch_size: int):
x = layers.Conv2D(filters, kernel_size=patch_size, strides=patch_size)(x)
return activation_block(x)
+2、ConvMixer 前的激活块
def activation_block(x):
x = layers.Activation('gelu')(x)
return layers.BatchNormalization()(x)
+3、ConvMixer 块
def conv_mixer_block(x, filters: int, kernel_size: int):
# Depthwise 卷积
x0 = x
x = layers.DepthwiseConv2D(kernel_size=kernel_size, padding='same')(x)
x = layers.Add()([activation_block(x), x0]) # 残差连接
# Pointwise 卷积
x = layers.Conv2D(filters, kernel_size=1)(x)
x = activation_block(x)
return x
+4、完整的网络
def get_conv_mixer_256_8(
image_size=32, filters=256, depth=8, kernel_size=5, patch_size=2, num_classes=10
):
"""ConvMixer-256/8: https://openreview.net/pdf?id=TVHS5Y4dNvM.
The hyperparameter values are taken from the paper.
"""
inputs = keras.Input((image_size, image_size, 3))
x = layers.Rescaling(scale=1.0 / 255)(inputs)
# 计算 patch 的嵌入
x = conv_stem(x, filters, patch_size)
# ConvMixer 块，depth 层
for _ in range(depth):
x = conv_mixer_block(x, filters, kernel_size)
# 输入分类器
x = layers.GlobalAvgPool2D()(x)
outputs = layers.Dense(num_classes, activation='softmax')(x)
return keras.Model(inputs, outputs)
本实验中使用的模型称为 ConvMixer-256/8，其中 256 表示通道数，8 表示深度。
7模型训练和评估
def run_experiment(model):
optimizer = tfa.optimizers.AdamW(
learning_rate=learning_rate, weight_decay=weight_decay
)
model.compile(
optimizer=optimizer,
loss='sparse_categorical_crossentropy',
metrics=['accuracy'],
)
checkpoint_filepath = './checkpoint'
checkpoint_callback = keras.callbacks.ModelCheckpoint(
checkpoint_filepath,
monitor='val_accuracy',
save_best_only=True,
save_weights_only=True,
)
history = model.fit(
train_dataset,
validation_data=val_dataset,
epochs=num_epochs,
callbacks=[checkpoint_callback],
)
model.load_weights(checkpoint_filepath)
_, accuracy = model.evaluate(test_dataset)
print(f'Test accuracy: {round(accuracy * 100, 2)}%')
return history, model
conv_mixer_model = get_conv_mixer_256_8()
history, conv_mixer_model = run_experiment(conv_mixer_model)
Epoch 1/10
352/352 [==============================] - 84s 150ms/step - loss: 1.2139 - accuracy: 0.5626 - val_loss: 3.4178 - val_accuracy: 0.1010
Epoch 2/10
352/352 [==============================] - 52s 147ms/step - loss: 0.7774 - accuracy: 0.7291 - val_loss: 0.8245 - val_accuracy: 0.7088
Epoch 3/10
352/352 [==============================] - 52s 147ms/step - loss: 0.5902 - accuracy: 0.7955 - val_loss: 0.5996 - val_accuracy: 0.7938
Epoch 4/10
352/352 [==============================] - 52s 147ms/step - loss: 0.4836 - accuracy: 0.8330 - val_loss: 0.5909 - val_accuracy: 0.7966
Epoch 5/10
352/352 [==============================] - 52s 147ms/step - loss: 0.4038 - accuracy: 0.8619 - val_loss: 0.5585 - val_accuracy: 0.8062
Epoch 6/10
352/352 [==============================] - 52s 147ms/step - loss: 0.3450 - accuracy: 0.8803 - val_loss: 0.5237 - val_accuracy: 0.8168
Epoch 7/10
352/352 [==============================] - 52s 147ms/step - loss: 0.3019 - accuracy: 0.8970 - val_loss: 0.5351 - val_accuracy: 0.8270
Epoch 8/10
352/352 [==============================] - 52s 148ms/step - loss: 0.2618 - accuracy: 0.9096 - val_loss: 0.5051 - val_accuracy: 0.8352
Epoch 9/10
352/352 [==============================] - 51s 146ms/step - loss: 0.2363 - accuracy: 0.9168 - val_loss: 0.5453 - val_accuracy: 0.8260
Epoch 10/10
352/352 [==============================] - 51s 146ms/step - loss: 0.2128 - accuracy: 0.9269 - val_loss: 0.5667 - val_accuracy: 0.8290
79/79 [==============================] - 3s 39ms/step - loss: 0.5329 - accuracy: 0.8271
Test accuracy: 82.71%
虽然训练集和验证集上的性能差异较大，但这点可以通过额外的正则化技术来拉近。从结果看，具有 80 万个参数的网络模型能够在 10 个 epochs 内达到 ~83% 的准确度是一个不错的结果。是不是感觉它还是值得进一步学习和研究一下的。
看一下论文中给出的在 ImageNet-1k 上的性能比较，
8可视化
我们可以可视化 patch 嵌入和学习到的卷积滤波器。这里，每个 patch 嵌入和中间 feature map 都具有相同数量的通道数，即 256。
def visualization_plot(weights, idx=1):
p_min, p_max = weights.min(), weights.max()
weights = (weights - p_min) / (p_max - p_min)
num_filters = 256
plt.figure(figsize=(8, 8))
for i in range(num_filters):
current_weight = weights[:, :, :, i]
if current_weight.shape[-1] == 1:
current_weight = current_weight.squeeze()
ax = plt.subplot(16, 16, idx)
ax.set_xticks([])
ax.set_yticks([])
plt.imshow(current_weight, cmap='coolwarm')
idx += 1
# 可视化 patch 嵌入
patch_embeddings = conv_mixer_model.layers[2].get_weights()[0]
visualization_plot(patch_embeddings)
即使我们没有训练网络收敛，我们也可以注意到不同的 kernel 具有不同的模式。有些有相似之处，而有些则截然不同。这些可视化对于更大的图像尺寸将更显着。
同样，我们也可以可视化学习到的卷积核。
for i, layer in enumerate(conv_mixer_model.layers):
if isinstance(layer, layers.DepthwiseConv2D):
if layer.get_config()['kernel_size'] == (5, 5):
print(i, layer)
idx = 26 # 靠近网络中间选择一层 depthwise conv 展示
kernel = conv_mixer_model.layers[idx].get_weights()[0]
kernel = np.expand_dims(kernel.squeeze(), axis=2)
visualization_plot(kernel)
5
12
19
26
33
40
47
54
总共 256 个 filter，可以看到 kernel 中的不同 filter 具有不同的局部跨度，并且这种模式可能会随着更多的训练而演变。
9附录 - 深度可分离卷积
Separable convolution 早在 Google 的 Xception 以及 MobileNet 论文中就提出来了，这里仅供初学卷积神经网络的童鞋阅读。
它的核心思想是将一个传统卷积运算分解为两步进行，分别为 Depthwise convolution 与 Pointwise convolution。可以认为它将传统卷积从通道内与通道间两个角度分离。
+传统卷积运算
先看一个动画，感受一下传统的卷积运算。
输入数据具有 3 个大小为 5 × 5 的通道，即 3 × 5 × 5，而卷积核也同样是 3 个通道，大小为 3 × 3，因此共有 3 × 3 × 3 个参数。这里只有一个卷积核，因此最终得到一张 feature map。
再来看一个有两个卷积核从而得到两张 feature map 的动画。
可以看到，这个卷积运算每次都涉及空间三个维度。
好了，看过动画应该就很清楚传统卷积是怎么运算的了，接下进入静态模式。
假设输入层为一个大小为 5 × 5 像素、三通道彩色图片。经过一个包含 4 个 filter 的卷积层，最终输出 4 个 feature Map。如果使用 padding='same' 来填充，则尺寸与输入层相同 5 × 5，如果没有则尺寸缩小为 3 × 3。这个过程可以用下图来可视化，
此时，卷积层共 4 个 filter，每个 filter 包含了 3 个 kernel，每个 kernel 的大小为 3 × 3。因此该卷积层的参数数量为4 × 3 × 3 × 3。
传统卷积运算的特点是将通道内和通道间同时卷积，一次性抓取特征的空间结构。
+Depthwise Convolution
还是上述例子，大小为 5 × 5 像素、三通道彩色图片首先经过第一次卷积运算，不同之处在于此次的卷积完全是在单个通道内进行，且 filter 的数量与上一层的 depth 相同。
因此，一个三通道的图像经过运算后生成了 3 个 feature map，如下图所示。
其中一个 filter 只包含一个大小为 3 × 3 的 kernel，卷积部分的参数个数为 3 × 3 × 3 。
Depthwise convolution 完成后的 feature map 数量与输入层的 depth（通道数）相同。
该卷积操作并没有利用不同通道在相同空间位置上的结构。因此还需要将这些 feature map 进行组合生成新的 feature map，即下面的 Pointwise convolution。
+Pointwise Convolution
Pointwise convolution 的卷积核尺寸为 1 × 1 × M，M 为上一层 feature map 的通道数。所以这里的卷积运算会将上一步的 map 在不同通道间进行加权，生成新的 feature map。
有几个 filter 就有几个 feature map，如下图所示，
由于采用的是 1×1 卷积的方式，此步中卷积涉及到的参数个数为 1 × 1 × 3 × 4。
经过 Pointwise convolution 之后，同样输出了 4 张 feature map，与传统卷积的输出维度相同。
我们把上面两种卷积放在一张图里，
+参数对比
回顾一下，传统卷积的参数个数为 4 × 3 × 3 × 3；而 Separable convolution 的参数个数为 3 × 3 × 3 + 1 × 1 × 3 × 4。
输入相同，输出也是 4 张 feature map，而 Separable convolution 的参数个数大大少于常规卷积。因此，假设参数量相同，那么采用 Separable convolution 的神经网络可以具有更深的层次。
这里通过一个简单例子介绍了 Depthwise 和 Pointwise 两个卷积运算以及与传统卷积运算的关系。如果之前对这些概念不了解的童鞋可以再回过头去看 ConvMixer 了。