PointNet++

论文:https://arxiv.org/abs/1706.02413(NIPS 2017)

code: https://github.com/charlesq34/pointnet2

1、改进

PointNet因为是只使用了MLP和max pooling,没有能力捕获局部结构,因此在细节处理和泛化到复杂场景上能力很有限。

  1. point-wise MLP,仅仅是对每个点表征,对局部结构信息整合能力太弱 –> PointNet++的改进:sampling和grouping整合局部邻域
  2. global feature直接由max pooling获得,无论是对分类还是对分割任务,都会造成巨大的信息损失 –> PointNet++的改进:hierarchical feature learning framework,通过多个set abstraction逐级降采样,获得不同规模不同层次的local-global feature
  3. 分割任务的全局特征global feature是直接复制与local feature拼接,生成discriminative feature能力有限 –> PointNet++的改进:分割任务设计了encoder-decoder结构,先降采样再上采样,使用skip connection将对应层的local-global feature拼接

2、方法

PointNet++的网络大体是encoder-decoder结构

encoder为降采样过程,通过多个set abstraction结构实现多层次的降采样,得到不同规模的point-wise feature,最后一个set abstraction输出可以认为是global feature。其中set abstraction由sampling,grouping,pointnet三个模块构成。

decoder根据分类和分割应用,又有所不同。分类任务decoder比较简单,不介绍了。分割任务decoder为上采样过程,通过反向插值和skip connection实现在上采样的同时,还能够获得local+global的point-wise feature,使得最终的表征能够discriminative(分辩能力)。

思考:

  1. PointNet++降采样过程是怎么实现的?/PointNet++是如何表征global feature的?(关注set abstraction, sampling layer, grouping layer, pointnet layer)
  2. PointNet++用于分割任务的上采样过程是怎么实现的?/PointNet++是如何表征用于分割任务的point-wise feature的?(关注反向插值,skip connection)

 🐖:上图中的 d 表示坐标空间维度, C 表示特征空间维度

2.1 encoder

在PointNet的基础上增加了hierarchical (层级)feature learning framework的结构。这种多层次的结构由set abstraction层组成。

在每一个层次的set abstraction,点集都会被处理和抽象,而产生一个规模更小的点集,可以理解成是一个降采样表征过程,可参考上图左半部分。

set abstraction由三个部分构成(代码贴在下面):

def pointnet_sa_module(xyz, points, npoint, radius, nsample, mlp, mlp2, group_all, is_training, bn_decay, scope, bn=True, pooling='max', knn=False, use_xyz=True, use_nchw=False):
    ''' PointNet Set Abstraction (SA) Module
        Input:
            xyz: (batch_size, ndataset, 3) TF tensor
            points: (batch_size, ndataset, channel) TF tensor
            npoint: int32 -- #points sampled in farthest point sampling
            radius: float32 -- search radius in local region
            nsample: int32 -- how many points in each local region
            mlp: list of int32 -- output size for MLP on each point
            mlp2: list of int32 -- output size for MLP on each region
            group_all: bool -- group all points into one PC if set true, OVERRIDE
                npoint, radius and nsample settings
            use_xyz: bool, if True concat XYZ with local point features, otherwise just use point features
            use_nchw: bool, if True, use NCHW data format for conv2d, which is usually faster than NHWC format
        Return:
            new_xyz: (batch_size, npoint, 3) TF tensor
            new_points: (batch_size, npoint, mlp[-1] or mlp2[-1]) TF tensor
            idx: (batch_size, npoint, nsample) int32 -- indices for local regions
    '''
    data_format = 'NCHW' if use_nchw else 'NHWC'
    with tf.variable_scope(scope) as sc:
        # Sample and Grouping
        if group_all:
            nsample = xyz.get_shape()[1].value
            new_xyz, new_points, idx, grouped_xyz = sample_and_group_all(xyz, points, use_xyz)
        else:
            new_xyz, new_points, idx, grouped_xyz = sample_and_group(npoint, radius, nsample, xyz, points, knn, use_xyz)
        # Point Feature Embedding
        if use_nchw: new_points = tf.transpose(new_points, [0,3,1,2])
        for i, num_out_channel in enumerate(mlp):
            new_points = tf_util.conv2d(new_points, num_out_channel, [1,1],
                                        padding='VALID', stride=[1,1],
                                        bn=bn, is_training=is_training,
                                        scope='conv%d'%(i), bn_decay=bn_decay,
                                        data_format=data_format) 
        if use_nchw: new_points = tf.transpose(new_points, [0,2,3,1])
        # Pooling in Local Regions
        if pooling=='max':
            new_points = tf.reduce_max(new_points, axis=[2], keep_dims=True, name='maxpool')
        elif pooling=='avg':
            new_points = tf.reduce_mean(new_points, axis=[2], keep_dims=True, name='avgpool')
        elif pooling=='weighted_avg':
            with tf.variable_scope('weighted_avg'):
                dists = tf.norm(grouped_xyz,axis=-1,ord=2,keep_dims=True)
                exp_dists = tf.exp(-dists * 5)
                weights = exp_dists/tf.reduce_sum(exp_dists,axis=2,keep_dims=True) # (batch_size, npoint, nsample, 1)
                new_points *= weights # (batch_size, npoint, nsample, mlp[-1])
                new_points = tf.reduce_sum(new_points, axis=2, keep_dims=True)
        elif pooling=='max_and_avg':
            max_points = tf.reduce_max(new_points, axis=[2], keep_dims=True, name='maxpool')
            avg_points = tf.reduce_mean(new_points, axis=[2], keep_dims=True, name='avgpool')
            new_points = tf.concat([avg_points, max_points], axis=-1)
        # [Optional] Further Processing 
        if mlp2 is not None:
            if use_nchw: new_points = tf.transpose(new_points, [0,3,1,2])
            for i, num_out_channel in enumerate(mlp2):
                new_points = tf_util.conv2d(new_points, num_out_channel, [1,1],
                                            padding='VALID', stride=[1,1],
                                            bn=bn, is_training=is_training,
                                            scope='conv_post_%d'%(i), bn_decay=bn_decay,
                                            data_format=data_format) 
            if use_nchw: new_points = tf.transpose(new_points, [0,2,3,1])
        new_points = tf.squeeze(new_points, [2]) # (batch_size, npoints, mlp2[-1])
        return new_xyz, new_points, idx

2.1.1 sampling layer

使用FPS(最远点采样)对点集进行降采样,将输入点集从规模 N1 降到更小的规模 N2 。FPS可以理解成是使得采样的各个点之间尽可能远,这种采样的好处是可以降采样结果会比较均匀。

FPS实现方式如下:随机选择一个点作为初始点作为已选择采样点,计算未选择采样点集中每个点与已选择采样点集之间的距离distance,将距离最大的那个点加入已选择采样点集,然后更新distance,一直循环迭代下去,直至获得了目标数量的采样点。

class FarthestSampler:
    def __init__(self):
        pass
    def _calc_distances(self, p0, points):
        return ((p0 - points) ** 2).sum(axis=1)
    def __call__(self, pts, k):
        farthest_pts = np.zeros((k, 3), dtype=np.float32)
        farthest_pts[0] = pts[np.random.randint(len(pts))]
        distances = self._calc_distances(farthest_pts[0], pts)
        for i in range(1, k):
            farthest_pts[i] = pts[np.argmax(distances)]
            distances = np.minimum(
                distances, self._calc_distances(farthest_pts[i], pts))
        return farthest_pts

输入规模为 B∗N∗(d+C) ,其中 B 表示batch size, N 表示点集中点的数量, d 表示点的坐标维度, C 表示点的其他特征(比如法向量等)维度。一般 d=3 , c=0

输出规模为 B∗N1∗(d+C) , N1<N ,因为这是一个降采样过程。

sampling和grouping具体实现是写在一个函数里的:

def sample_and_group(npoint, radius, nsample, xyz, points, knn=False, use_xyz=True):
    '''
    Input:
        npoint: int32
        radius: float32
        nsample: int32
        xyz: (batch_size, ndataset, 3) TF tensor
        points: (batch_size, ndataset, channel) TF tensor, if None will just use xyz as points
        knn: bool, if True use kNN instead of radius search
        use_xyz: bool, if True concat XYZ with local point features, otherwise just use point features
    Output:
        new_xyz: (batch_size, npoint, 3) TF tensor
        new_points: (batch_size, npoint, nsample, 3+channel) TF tensor
        idx: (batch_size, npoint, nsample) TF tensor, indices of local points as in ndataset points
        grouped_xyz: (batch_size, npoint, nsample, 3) TF tensor, normalized point XYZs
            (subtracted by seed point XYZ) in local regions
    '''
    new_xyz = gather_point(xyz, farthest_point_sample(npoint, xyz)) # (batch_size, npoint, 3)
    if knn:
        _,idx = knn_point(nsample, xyz, new_xyz)
    else:
        idx, pts_cnt = query_ball_point(radius, nsample, xyz, new_xyz)
    grouped_xyz = group_point(xyz, idx) # (batch_size, npoint, nsample, 3)
    grouped_xyz -= tf.tile(tf.expand_dims(new_xyz, 2), [1,1,nsample,1]) # translation normalization
    if points is not None:
        grouped_points = group_point(points, idx) # (batch_size, npoint, nsample, channel)
        if use_xyz:
            new_points = tf.concat([grouped_xyz, grouped_points], axis=-1) # (batch_size, npoint, nample, 3+channel)
        else:
            new_points = grouped_points
    else:
        new_points = grouped_xyz
    return new_xyz, new_points, idx, grouped_xyz

其中sampling对应的部分是:

new_xyz = gather_point(xyz, farthest_point_sample(npoint, xyz)) # (batch_size, npoint, 3)

xyz既是 B∗N∗3 的点云,npoint是降采样点的规模。注意:PointNet++的FPS均是在坐标空间做的,而不是在特征空间做的。这一点很关键,因为FPS本身是不可微的,无法计算梯度反向传播。

本着刨根问题的心态,我们来看看farthest_point_sample和gather_point究竟在做什么

farthest_point_sample输入输出非常明晰,输出的是降采样点在inp中的索引,因此是 B∗N1 int32类型的张量

def farthest_point_sample(npoint,inp):
    '''
input:
    int32
    batch_size * ndataset * 3   float32
returns:
    batch_size * npoint         int32
    '''
    return sampling_module.farthest_point_sample(inp, npoint)

gather_point的作用就是将上面输出的索引,转化成真正的点云

def gather_point(inp,idx):
    '''
input:
    batch_size * ndataset * 3   float32
    batch_size * npoints        int32
returns:
    batch_size * npoints * 3    float32
    '''
    return sampling_module.gather_point(inp,idx)

2.1.2 grouping layer

上一步sampling的过程是将 N∗(d+C) 降到 N1∗(d+C) (这里论述方便先不考虑batch,就考虑单个点云),实际上可以理解成是在 N 个点中选取 N1 个中心点(key point)。

那么这一步grouping的目的就是以这每个key point为中心,找其固定规模(令规模为 K)的邻点,共同组成一个局部邻域(patch)。也就是会生成 N1 个局部邻域,输出规模为 N1∗K∗(d+C)

if knn:
    _,idx = knn_point(nsample, xyz, new_xyz)
else:
    idx, pts_cnt = query_ball_point(radius, nsample, xyz, new_xyz)
    grouped_xyz = group_point(xyz, idx) # (batch_size, npoint, nsample, 3)

1)找邻域的过程也是在坐标空间进行(也就是以上代码输入输出维度都是 d ,没有 C  C 是在后面的代码拼接上的),而不是特征空间。

2)找邻域这里有两种方式:KNN和query ball point.

其中前者KNN就是大家耳熟能详的K近邻,找K个坐标空间最近的点。
后者query ball point就是划定某一半径,找在该半径球内的点作为邻点。

还有个问题:query ball point如何保证对于每个局部邻域,采样点的数量都是一样的呢?
事实上,如果query ball的点数量大于规模 K ,那么直接取前 K 个作为局部邻域;如果小于,那么直接对某个点重采样,凑够规模 K

KNN和query ball的区别:(摘自原文)Compared with kNN, ball query’s local neighborhood guarantees a fixed region scale thus making local region feature more generalizable across space, which is preferred for tasks requiring local pattern recognition (e.g. semantic point labeling).也就是query ball更加适合于应用在局部/细节识别的应用上,比如局部分割。

补充材料中也有实验来对比KNN和query ball:

sample_and_group代码的剩余部分:

sample和group操作都是在坐标空间进行的,因此如果还有特征空间信息(即point-wise feature),可以在这里将其与坐标空间拼接,组成新的point-wise feature,准备送入后面的unit point进行特征学习。

if points is not None:
    grouped_points = group_point(points, idx) # (batch_size, npoint, nsample, channel)
    if use_xyz:
        new_points = tf.concat([grouped_xyz, grouped_points], axis=-1) # (batch_size, npoint, nample, 3+channel)
    else:
        new_points = grouped_points
else:
    new_points = grouped_xyz

2.1.3 PointNet layer

使用PointNet对以上结果表征

输入 B∗N∗K∗(d+C) ,输出 B∗N∗(d+C1)

以下代码主要分成3个部分:

1)point feature embedding

这里输入是 B∗N∗K∗(d+C) ,可以类比成是batch size为 B ,宽高为 N∗K ,通道数为 d+C 的图像,这样一类比,这里的卷积就好理解多了实际上就是 1∗1 卷积,不改变feature map大小,只改变通道数,将通道数升高,实现所谓“embedding”

这部分输出是 B∗N∗K∗C1

2)pooling in local regions

pooling,只是是对每个局部邻域pooling,输出是 B∗N∗1∗C1

3)further processing

再对池化后的结果做MLP,也是简单的 1∗1 卷积。这一部分在实际实验中PointNet++并没有设置去做

# Point Feature Embedding
if use_nchw: new_points = tf.transpose(new_points, [0,3,1,2])
for i, num_out_channel in enumerate(mlp):
    new_points = tf_util.conv2d(new_points, num_out_channel, [1,1],
                                padding='VALID', stride=[1,1],
                                bn=bn, is_training=is_training,
                                scope='conv%d'%(i), bn_decay=bn_decay,
                                data_format=data_format) 
if use_nchw: new_points = tf.transpose(new_points, [0,2,3,1])

# Pooling in Local Regions
if pooling=='max':
    new_points = tf.reduce_max(new_points, axis=[2], keep_dims=True, name='maxpool')
elif pooling=='avg':
    new_points = tf.reduce_mean(new_points, axis=[2], keep_dims=True, name='avgpool')
elif pooling=='weighted_avg':
    with tf.variable_scope('weighted_avg'):
        dists = tf.norm(grouped_xyz,axis=-1,ord=2,keep_dims=True)
        exp_dists = tf.exp(-dists * 5)
        weights = exp_dists/tf.reduce_sum(exp_dists,axis=2,keep_dims=True) # (batch_size, npoint, nsample, 1)
        new_points *= weights # (batch_size, npoint, nsample, mlp[-1])
        new_points = tf.reduce_sum(new_points, axis=2, keep_dims=True)
elif pooling=='max_and_avg':
    max_points = tf.reduce_max(new_points, axis=[2], keep_dims=True, name='maxpool')
    avg_points = tf.reduce_mean(new_points, axis=[2], keep_dims=True, name='avgpool')
    new_points = tf.concat([avg_points, max_points], axis=-1)

# [Optional] Further Processing 
if mlp2 is not None:
    if use_nchw: new_points = tf.transpose(new_points, [0,3,1,2])
    for i, num_out_channel in enumerate(mlp2):
        new_points = tf_util.conv2d(new_points, num_out_channel, [1,1],
                                    padding='VALID', stride=[1,1],
                                    bn=bn, is_training=is_training,
                                    scope='conv_post_%d'%(i), bn_decay=bn_decay,
                                    data_format=data_format) 
    if use_nchw: new_points = tf.transpose(new_points, [0,2,3,1])

2.1.4 Encoder还有一个问题

pointnet++实际上就是对局部邻域表征。

那就不得不面对一个挑战:non-uniform sampling density(点云的密度不均匀),也就是在稀疏点云局部邻域训练可能不能很好挖掘点云的局部结构

PointNet++做法:learn to combine features from regions of different scales when the input sampling density changes.

因此文章提出了两个方案:

一、Multi-scale grouping(MSG)

对当前层的每个中心点,取不同radius的query ball,可以得到多个不同大小的同心球,也就是得到了多个相同中心但规模不同的局部邻域,分别对这些局部邻域表征,并将所有表征拼接。如上图所示。

该方法比较麻烦,运算较多。

代码层面其实就是加了个遍历radius_list的循环,分别处理,并最后concat

new_xyz = gather_point(xyz, farthest_point_sample(npoint, xyz))
new_points_list = []
for i in range(len(radius_list)):
    radius = radius_list[i]
    nsample = nsample_list[i]
    idx, pts_cnt = query_ball_point(radius, nsample, xyz, new_xyz)
    grouped_xyz = group_point(xyz, idx)
    grouped_xyz -= tf.tile(tf.expand_dims(new_xyz, 2), [1,1,nsample,1])
    if points is not None:
        grouped_points = group_point(points, idx)
        if use_xyz:
            grouped_points = tf.concat([grouped_points, grouped_xyz], axis=-1)
    else:
        grouped_points = grouped_xyz
    if use_nchw: grouped_points = tf.transpose(grouped_points, [0,3,1,2])
    for j,num_out_channel in enumerate(mlp_list[i]):
        grouped_points = tf_util.conv2d(grouped_points, num_out_channel, [1,1],
                                        padding='VALID', stride=[1,1], bn=bn, is_training=is_training,
                                        scope='conv%d_%d'%(i,j), bn_decay=bn_decay)
    if use_nchw: grouped_points = tf.transpose(grouped_points, [0,2,3,1])
    new_points = tf.reduce_max(grouped_points, axis=[2])
    new_points_list.append(new_points)
new_points_concat = tf.concat(new_points_list, axis=-1)

二、Multi-resolution grouping(MRG)

(摘自原文)features of a region at some level Li is a concatenation of two vectors.

One vector (left in figure) is obtained by summarizing the features at each subregion from the lower level Li−1 using the set abstraction level.

The other vector (right) is the feature that is obtained by directly processing all raw points in the local region using a single PointNet.

简单来说,就是当前set abstraction的局部邻域表征由两部分构成:

左边表征:对上一层set abstraction(还记得上一层的点规模是更大的吗?)各个局部邻域(或者说中心点)的特征进行聚合。 右边表征:使用一个单一的PointNet直接在局部邻域处理原始点云

2.2 decoder:

2.2.1 分类任务的decoder

比较简单,将encoder降采样得到的global feature送入几层全连接网络,最后通过一个softmax分类。

2.2.2 分割任务的decoder

经过前半部分的encoder,我们得到的是global feature,或者是极少数点的表征(其实也就是global feature)

而如果做分割,我们需要的是point-wise feature,这可怎么办呢?

PointNet处理思路很简单,直接把global feature复制并与之前的local feature拼接,使得这个新point-wise feature能够获得一定程度的“邻域”信息。这种简单粗暴的方法显然并不能得到很discriminative的表征

别急,PointNet++来了。

PointNet++设计了一种反向插值的方法来实现上采样的decoder结构,通过反向插值和skip connection来获得discriminative point-wise feature:

设红色矩形点集 P1 : N1∗C ,蓝色矩形点集 P2 : N2∗C2 ,因为decoder是上采样过程,因此 N2>N1

一、反向插值具体做法:

对于 P2 中的每个点 x ,找在原始点云坐标空间下, P1 中与其最接近的 k 个点 x1,…,xk

当前我们想通过反向插值的方式用较少的点把更多的点的特征插出来,实现上采样

此时 x1,…,xk 的特征我们是知道的,我们想得到 x 的特征。如上公式,实际上就是将 x1,…,xk 的特征加权求和,得到x的特征。其中这个权重是与x和 x1,…,xk 的距离成反向相关的,意思就是距离越远的点,对x特征的贡献程度越小。P2 中其他点以此类推,从而实现了特征的上采样回传

skip connection具体做法:

回传得到的point-wise feature是从decoder的上一层得到的,因此算是global级别的信息,这对于想得到discriminative还是不够,因为我们还缺少local级别的信息!!!

如上图就是我们反向插值只得到了 C2 ,但是我们还需要提供local级别信息的 C1 特征!!!

这时skip connection来了!!!

skip connection其实就是将之前encoder对应层的表征直接拼接了过来

因为上图中encoder蓝色矩形点集的 C1 表征是来自于规模更大的绿色矩形点集的表征,这在一定程度上其实是实现了local级别的信息

我们通过反向插值和skip connection在decoder中逐级上采样得到local + global point-wise feature,得到了discriminative feature,应用于分割任务。

2.3 loss

无论是分类还是分割应用,本质上都是分类问题,因此loss就是分类任务中常用的交叉熵loss

2.4 其他的问题

Q:PointNet++梯度是如何回传的???

A:PointNet++ fps实际上并没有参与梯度计算和反向传播。

可以理解成是PointNet++将点云进行不同规模的fps降采样,事先将这些数据准备好,再送到网络中去训练的

3 dataset and experiments

3.1 dataset

  • MNIST: Images of handwritten digits with 60k training and 10k testing samples.(用于分类)
  • ModelNet40: CAD models of 40 categories (mostly man-made). We use the official split with 9,843 shapes for training and 2,468 for testing. (用于分类)
  • SHREC15: 1200 shapes from 50 categories. Each category contains 24 shapes which are mostly organic ones with various poses such as horses, cats, etc. We use five fold cross validation to acquire classification accuracy on this dataset. (用于分类)
  • ScanNet: 1513 scanned and reconstructed indoor scenes. We follow the experiment setting in [5] and use 1201 scenes for training, 312 scenes for test. (用于分割)

3.2 experiments

主要关心的实验结果是2个:

  1. ModelNet40分类结果
  2. ShapeNet Part分割结果
PointNet++也做了ShapeNet part数据集上的part segmentation:

4、 conclusion

PointNet++是PointNet的续作,在一定程度上弥补了PointNet的一些缺陷,表征网络基本和PN类似,还是MLP、 1∗1 卷积、pooling那一套,核心创新点在于设计了局部邻域的采样表征方法和这种多层次的encoder-decoder结合的网络结构。

第一次看到PointNet++网络结构,觉得设计得非常精妙,特别是设计了上采样和下采样的具体实现方法,并以此用于分割任务的表征,觉得设计得太漂亮了。但其实无论是分类还是分割任务,提升幅度较PointNet也就是1-2个点而已。

PointNet++,特别是其前半部分encoder,提供了非常好的表征网络,后面很多点云处理应用的论文都会使用到PointNet++作为它们的表征器。

pointnet–基于点云的分类和分割深度学习算法

论文:https://arxiv.org/abs/1612.00593(cvpr2017)

code:https://github.com/charlesq34/pointnet

基础知识:

1、什么是点云?

简单来说就是一堆三维点的集合,必须包括各个点的三维坐标信息,其他信息比如各个点的法向量、颜色等均是可选。点云的文件格式可以有很多种,包括xyz,npy,ply,obj,off等(有些是mesh不过问题不大,因为mesh可以通过泊松采样等方式转化成点云)。对于单个点云,如果你使用np.loadtxt得到的实际上就是一个维度为 (num_points,num_channels) 的张量,num_channels一般为3,表示点云的三维坐标。

这里以horse.xyz文件为例,实际就是文本文件,打开后数据长这样(局部,总共有2048个点):

实际就是一堆点的信息,这里只有三维坐标,将其可视化出来长这样:

2、点云处理任务是重要的

三维图形具有多种表现形式,包括了mesh、体素、点云等,甚至还有些方法使用多视图来对三维图形表征。而点云在以上各种形式的数据中算是日常生活中最能够大规模获取和使用的数据结构了,包括自动驾驶、增强现实等在内的应用需要直接或间接从点云中提取信息,点云处理也逐渐成为计算机视觉非常重要的一部分。

正文:

PointNet所作的事情就是对点云做特征学习,并将学习到的特征去做不同的应用:分类(shape-wise feature)、分割(point-wise feature)等。

PointNet之所以影响力巨大,就是因为它为点云处理提供了一个简单、高效、强大的特征提取器(encoder),几乎可以应用到点云处理的各个应用中,其地位类似于图像领域的AlexNet。

1、动机

点云或者mesh,大多数研究人员都是将其转化成3D体素或者多视图来做特征学习的,这其中的工作包括了VoxelNet, MVCNN等。这些工作都或多或少存在了一些问题。

直接对点云做特征学习也不是不可以,但有几个问题需要考虑:特征学习需要对点云中各个点的排列保持不变性、特征学习需要对rigid transformation保持不变性等。虽然有挑战,但是深度学习强大的表征能力以及其在图像领域取得的巨大成功,因此是很有必要直接在点云上进行尝试的。

2、贡献

  1. 我们设计了一个新颖的深层网络架构来处理三维中的无序点集
  2. 我们设计的网络表征可以做三维图形分类、图形的局部分割以及场景的语义分割等任务
  3. 我们提供了完备的经验和理论分析来证明PointNet的稳定和高效。
  4. 充分的消融实验,证明网络各个部分对于表征的有效性。

3、方法

3.1 点云的几个特点:

  1. 无序性 –> 对称函数设计用于表征
  2. 点不是孤立的,需要考虑局部结构 –> 局部全局特征结合
  3. 仿射变换无关性 –> alignment network

(重要)关于第三点:相同的点云在空间中经过一定的刚性变化(旋转或平移),坐标发生变化。其实对于点云分类or分割任务来说(分割可以认为是点的分类),例如,整体的旋转和平移不应修改全局点云类别和每个点的类别,也不应修改点的分割因此需要保证仿射变换无关性(简单来说,“仿射变换”就是:“线性变换”+“平移”),但是对于位置 敏感的 任务:点云配准、点云补全任务,对于位置敏感,就不需要保证 仿射变换的无关性

我们希望不论点云在怎样的坐标系下呈现,网络都能正确的识别出。这个问题可以通过STN(spacial transform netw)来解决。三维不太一样的是点云是一个不规则的结构(无序,无网格),不需要重采样的过程。pointnet通过学习一个矩阵来达到对目标最有效的变换。

解决方法

  1. 空间变换网络解决旋转问题:三维的STN可以通过学习点云本身的位姿信息学习到一个最有利于网络进行分类或分割的DxD旋转矩阵(D代表特征维度,pointnet中D采用3和64)。至于其中的原理,我的理解是,通过控制最后的loss来对变换矩阵进行调整,pointnet并不关心最后真正做了什么变换,只要有利于最后的结果都可以。pointnet采用了两次STN,第一次input transform是对空间中点云进行调整,直观上理解是旋转出一个更有利于分类或分割的角度,比如把物体转到正面;第二次feature transform是对提取出的64维特征进行对齐,即在特征层面对点云进行变换。
  2. maxpooling解决无序性问题:网络对每个点进行了一定程度的特征提取之后,maxpooling可以对点云的整体提取出global feature。

3.2 网络结构:

batchnormal对于上采样任务来说效果不好

网络分成了分类网络和分割网络2个部分,大体思路类似,都是设计表征的过程分类网络设计global feature,分割网络设计point-wise feature。两者都是为了让表征尽可能discriminative,也就是同类的能分到一类,不同类的距离能拉开。

输入 n*3 n是点数

inputtransform:放射变换(为了保证仿射变换的不变性):直接预测一个变换矩阵(3*3)来处理输入点的坐标(对所有坐标进行变换)。因为会有数据增强的操作存在,这样做可以在一定程度上保证网络可以学习到变换无关性。T-Net模型,它的主要作用是学习出变化矩阵来对输入的点云或特征进行规范化处理。

MLP:

有两种实现方法:

1、输入 B,N,3 —- nn.liner层 — B,N,64

2、输入 B,3,N —- conv1d(1×1) — B,64,N

Pooling:

为了解决无序性(点云本质上是一长串点(nx3矩阵,其中n是点数)。在几何上,点的顺序不影响它在空间中对整体形状的表示,例如,相同的点云可以由两个完全不同的矩阵表示。)使用 maxpooling或sumpooling,也就是说,最后的D维特征对每一维都选取N个点中对应的最大特征值或特征值总和,这样就可以通过g来解决无序性问题。

最后再经过一个mlp(代码中运用全连接)得到k个score。分类网络最后接的loss是softmax。

分割网络:

将池化后的特征和前一阶段特征拼接,池化后的特征有全局信息,在和之前的拼接,以此得到同时对局部信息和全局信息感知的point-wise特征,提升表征效果。然后最后输出n*m, m为类别数量,表示每个点的类别信息。

损失函数:

分类中常用的交叉熵+alignment network中用于约束生成的alignment matrix的loss

dataset and experiments

evaluate metric

分类:分类准确率acc
分割:mIoU

dataset

分类:ModelNet40
分割:ShapeNet Part dataset和Stanford 3D semantic parsing dataset

experiments

1、分类:

2、局部分割:

code:

1. 如何对点云使用MLP?
2. alignment network怎么做的?
3. 对称函数如何实现来提取global feature的?
4. loss?

def get_model(point_cloud, is_training, bn_decay=None):
    """ Classification PointNet, input is BxNx3, output Bx40 """
    batch_size = point_cloud.get_shape()[0].value
    num_point = point_cloud.get_shape()[1].value
    end_points = {}
    with tf.variable_scope('transform_net1') as sc:
        transform = input_transform_net(point_cloud, is_training, bn_decay, K=3)
    point_cloud_transformed = tf.matmul(point_cloud, transform)
    input_image = tf.expand_dims(point_cloud_transformed, -1)
    net = tf_util.conv2d(input_image, 64, [1,3],
                         padding='VALID', stride=[1,1],
                         bn=True, is_training=is_training,
                         scope='conv1', bn_decay=bn_decay)
    net = tf_util.conv2d(net, 64, [1,1],
                         padding='VALID', stride=[1,1],
                         bn=True, is_training=is_training,
                         scope='conv2', bn_decay=bn_decay)
    with tf.variable_scope('transform_net2') as sc:
        transform = feature_transform_net(net, is_training, bn_decay, K=64)
    end_points['transform'] = transform
    net_transformed = tf.matmul(tf.squeeze(net, axis=[2]), transform)
    net_transformed = tf.expand_dims(net_transformed, [2])
    net = tf_util.conv2d(net_transformed, 64, [1,1],
                         padding='VALID', stride=[1,1],
                         bn=True, is_training=is_training,
                         scope='conv3', bn_decay=bn_decay)
    net = tf_util.conv2d(net, 128, [1,1],
                         padding='VALID', stride=[1,1],
                         bn=True, is_training=is_training,
                         scope='conv4', bn_decay=bn_decay)
    net = tf_util.conv2d(net, 1024, [1,1],
                         padding='VALID', stride=[1,1],
                         bn=True, is_training=is_training,
                         scope='conv5', bn_decay=bn_decay)
    # Symmetric function: max pooling
    net = tf_util.max_pool2d(net, [num_point,1],
                             padding='VALID', scope='maxpool')
    net = tf.reshape(net, [batch_size, -1])
    net = tf_util.fully_connected(net, 512, bn=True, is_training=is_training,
                                  scope='fc1', bn_decay=bn_decay)
    net = tf_util.dropout(net, keep_prob=0.7, is_training=is_training,
                          scope='dp1')
    net = tf_util.fully_connected(net, 256, bn=True, is_training=is_training,
                                  scope='fc2', bn_decay=bn_decay)
    net = tf_util.dropout(net, keep_prob=0.7, is_training=is_training,
                          scope='dp2')
    net = tf_util.fully_connected(net, 40, activation_fn=None, scope='fc3')
    return net, end_points

MLP的核心做法:

input_image = tf.expand_dims(point_cloud_transformed, -1)
net = tf_util.conv2d(input_image, 64, [1,3],
                         padding='VALID', stride=[1,1],
                         bn=True, is_training=is_training,
                         scope='conv1', bn_decay=bn_decay)
net = tf_util.conv2d(net, 64, [1,1],
                     padding='VALID', stride=[1,1],
                     bn=True, is_training=is_training,
                     scope='conv2', bn_decay=bn_decay)

这里input_image维度是 B×N×3×1 ,因此将点云看成是W和H分为N和3的2D图像,维度是 1

然后直接基于这个“2D图像”做卷积,第一个卷积核size是 [1,3] ,正好对应的就是“2D图像”的一行,也就是一个点(三维坐标),输出通道数是64,因此输出张量维度应该是 B×N×1×64

第二个卷积核size是 [1,1] , 1∗1 卷积只改变通道数,输出张量维度是 B×N×1×64

conv2d就是将卷积封装了一下,核心部分也就是调用tf.nn.conv2d,实现如下:

def conv2d(inputs,
           num_output_channels,
           kernel_size,
           scope,
           stride=[1, 1],
           padding='SAME',
           use_xavier=True,
           stddev=1e-3,
           weight_decay=0.0,
           activation_fn=tf.nn.relu,
           bn=False,
           bn_decay=None,
           is_training=None):
  """ 2D convolution with non-linear operation.
  Args:
    inputs: 4-D tensor variable BxHxWxC
    num_output_channels: int
    kernel_size: a list of 2 ints
    scope: string
    stride: a list of 2 ints
    padding: 'SAME' or 'VALID'
    use_xavier: bool, use xavier_initializer if true
    stddev: float, stddev for truncated_normal init
    weight_decay: float
    activation_fn: function
    bn: bool, whether to use batch norm
    bn_decay: float or float tensor variable in [0,1]
    is_training: bool Tensor variable
  Returns:
    Variable tensor
  """
  with tf.variable_scope(scope) as sc:
      kernel_h, kernel_w = kernel_size
      num_in_channels = inputs.get_shape()[-1].value
      kernel_shape = [kernel_h, kernel_w,
                      num_in_channels, num_output_channels]
      kernel = _variable_with_weight_decay('weights',
                                           shape=kernel_shape,
                                           use_xavier=use_xavier,
                                           stddev=stddev,
                                           wd=weight_decay)
      stride_h, stride_w = stride
      outputs = tf.nn.conv2d(inputs, kernel,
                             [1, stride_h, stride_w, 1],
                             padding=padding)
      biases = _variable_on_cpu('biases', [num_output_channels],
                                tf.constant_initializer(0.0))
      outputs = tf.nn.bias_add(outputs, biases)
      if bn:
        outputs = batch_norm_for_conv2d(outputs, is_training,
                                        bn_decay=bn_decay, scope='bn')
      if activation_fn is not None:
        outputs = activation_fn(outputs)
      return outputs

alignment network :

input_transform_net为例:

def input_transform_net(point_cloud, is_training, bn_decay=None, K=3):
    """ Input (XYZ) Transform Net, input is BxNx3 gray image
        Return:
            Transformation matrix of size 3xK """
    batch_size = point_cloud.get_shape()[0].value
    num_point = point_cloud.get_shape()[1].value
    input_image = tf.expand_dims(point_cloud, -1)
    net = tf_util.conv2d(input_image, 64, [1,3],
                         padding='VALID', stride=[1,1],
                         bn=True, is_training=is_training,
                         scope='tconv1', bn_decay=bn_decay)
    net = tf_util.conv2d(net, 128, [1,1],
                         padding='VALID', stride=[1,1],
                         bn=True, is_training=is_training,
                         scope='tconv2', bn_decay=bn_decay)
    net = tf_util.conv2d(net, 1024, [1,1],
                         padding='VALID', stride=[1,1],
                         bn=True, is_training=is_training,
                         scope='tconv3', bn_decay=bn_decay)
    net = tf_util.max_pool2d(net, [num_point,1],
                             padding='VALID', scope='tmaxpool')
    net = tf.reshape(net, [batch_size, -1])
    net = tf_util.fully_connected(net, 512, bn=True, is_training=is_training,
                                  scope='tfc1', bn_decay=bn_decay)
    net = tf_util.fully_connected(net, 256, bn=True, is_training=is_training,
                                  scope='tfc2', bn_decay=bn_decay)
    with tf.variable_scope('transform_XYZ') as sc:
        assert(K==3)
        weights = tf.get_variable('weights', [256, 3*K],
                                  initializer=tf.constant_initializer(0.0),
                                  dtype=tf.float32)
        biases = tf.get_variable('biases', [3*K],
                                 initializer=tf.constant_initializer(0.0),
                                 dtype=tf.float32)
        biases += tf.constant([1,0,0,0,1,0,0,0,1], dtype=tf.float32)
        transform = tf.matmul(net, weights)
        transform = tf.nn.bias_add(transform, biases)
    transform = tf.reshape(transform, [batch_size, 3, K])
    return transform

实际上,前半部分就是通过卷积和max_pooling对batch内各个点云提取global feature,再将global feature降到 3×K 维度,并reshape成 3×3 ,得到transform matrix

通过数据增强丰富训练数据集,网络确实应该学习到有效的transform matrix,用来实现transformation invariance

loss

监督分类任务中常用的交叉熵loss + alignment network中的mat_diff_loss

 对于特征空间的alignment network,由于特征空间维度比较高,因此直接生成的alignment matrix会维度特别大,不好优化,因此这里需要加个loss约束一下。

总结:

PointNet之所以影响力巨大,并不仅仅是因为它是第一篇,更重要的是它的网络很简洁(简洁中蕴含了大量的工作来探寻出简洁这条路)却非常的work,这也就使得它能够成为一个工具,一个为点云表征的encoder工具,应用到更广阔的点云处理任务中。

MLP+max pooling竟然就击败了众多SOTA,令人惊讶。另外PointNet在众多细节设计也都进行了理论分析和消融实验验证,保证了严谨性,这也为PointNet后面能够大规模被应用提供了支持。

让网络来学习resize:插即用的新型图像调整器模型

Learning to Resize Images for Computer Vision Tasks

论文地址:https://arxiv.org/abs/2103.09950

代码:https://github.com/KushajveerSingh/resize_network_cv

尽管近年来卷积神经网络很大地促进了计算机视觉的发展,但一个重要方面很少被关注:图像大小对被训练的任务的准确性的影响 。通常,输入图像的大小被调整到一个相对较小的空间分辨率(例如,224×224),然后再进行训练和推理。这种调整大小的机制通常是固定的图像调整器(image resizer)(如:双行线插值)但是这些调整器是否限制了训练网络的任务性能呢? 作者通过实验证明了典型的线性调整器可以被可学习的调整器取代,从而大大提高性能 。虽然经典的调整器通常会具备更好的小图像感知质量(即对人类识别图片更加友好),本文提出的可学习调整器不一定会具备更好的视觉质量,但能够提高CV任务的性能。

在不同的任务中,可学习的图像调整器与baseline视觉模型进行联合训练。这种可学习的基于cnn的调整器创建了机器友好的视觉操作,因此在不同的视觉任务中表现出了更好的性能 。作者使用ImageNet数据集来进行分类任务,实验中使用四种不同的baseline模型来学习不同的调整器,相比于baseline模型,使用本文提出的可学习调整器能够获得更高的性能提升。

主要包括了两个重要的特性:(1) 双线性特征调整大小(bilinear feature resizing),以及(2)跳过连接(skip connection),该连接可容纳双线性调整大小的图像和CNN功能的组合。

第一个特性考虑到以原始分辨率计算的特征与模型的一致性。跳过连接可以简化学习过程,因为重定大小器模型可以直接将双线性重定大小的图像传递到基线任务中。

与一般的编码器-解码器架构不同,这篇论文中所提出的体系结构允许将图像大小调整为任何目标大小和纵横比(注意:这个大小必须是我们自己设定的,而不是网络 自己学习的)。并且可学习的resizer性能几乎不依赖于双线性重定器的选择,这意味着它可以直接替换其他现成的方法。

以上之后,就没有别的了,还以为是什么样子的惊天设计,最后不就是:给网络变得复杂了吗,把这种复杂说成是可学习的resizer,这样的话,普通网络的浅层都可以说成是可学习的resizer不是吗?

另外,通过一些实验 来说,确实能够提升效果。个人认为 作者提出的resizer模型实际上是一个可训练的数据增强方法,甚至也可以认为就是将模型变得更加复杂。整体的网络就像是一般模型中resblock。

作者的对比试验是这样做的:首先通过常用的reisze方法训练网络模型,作为baseline,然后在训练好的网络模型前面添加可学习的resizer,然后进行训练,作为自己的方法。感受一下作者的实验结果吧。如表3和表4

表3 训练配置
表4 对比试验

四. 总结

文章写得好,加上点运气,都可以发高质量论文,以上

代码实现:

import torch
import torch.nn as nn
import torch.nn.functional as F
from functools import partial

"""
    Learning to Resize Images for Computer Vision Tasks
    https://arxiv.org/pdf/2105.04714.pdf
"""

def conv1x1(in_chs, out_chs = 16):
    return nn.Conv2d(in_chs, out_chs, kernel_size=1, stride=1, padding=0)


def conv3x3(in_chs, out_chs = 16):
    return nn.Conv2d(in_chs, out_chs, kernel_size=3, stride=1, padding=1)


def conv7x7(in_chs, out_chs = 16):
    return nn.Conv2d(in_chs, out_chs, kernel_size=7, stride=1, padding=3)


class ResBlock(nn.Module):
    def __init__(self, in_chs,out_chs = 16):
        super(ResBlock, self).__init__()
        self.layers = nn.Sequential(
            conv3x3(in_chs, out_chs),
            nn.BatchNorm2d(out_chs),
            nn.LeakyReLU(0.2),
            conv3x3(out_chs, out_chs),
            nn.BatchNorm2d(out_chs)
        )
    def forward(self, x):
        identity = x
        out = self.layers(x)
        out += identity
        return out


class Resizer(nn.Module):
    def __init__(self, in_chs, out_size, n_filters = 16, n_res_blocks = 1, mode = 'bilinear'):
        super(Resizer, self).__init__()
        self.interpolate_layer = partial(F.interpolate, size=out_size, mode=mode,
            align_corners=(True if mode in ('linear', 'bilinear', 'bicubic', 'trilinear') else None))
        self.conv_layers = nn.Sequential(
            conv7x7(in_chs, n_filters),
            nn.LeakyReLU(0.2),
            conv1x1(n_filters, n_filters),
            nn.LeakyReLU(0.2),
            nn.BatchNorm2d(n_filters)
        )
        self.residual_layers = nn.Sequential()
        for i in range(n_res_blocks):
            self.residual_layers.add_module(f'res{i}', ResBlock(n_filters, n_filters))
        self.residual_layers.add_module('conv3x3', conv3x3(n_filters, n_filters))
        self.residual_layers.add_module('bn', nn.BatchNorm2d(n_filters))
        self.final_conv = conv7x7(n_filters, in_chs)

    def forward(self, x):
        identity = self.interpolate_layer(x)
        conv_out = self.conv_layers(x)
        conv_out = self.interpolate_layer(conv_out)
        conv_out_identity = conv_out
        res_out = self.residual_layers(conv_out)
        res_out += conv_out_identity
        out = self.final_conv(res_out)
        out += identity
        return

Dynamic ReLU(2020)

论文地址: https://arxiv.org/pdf/2003.10027.pdf
源码地址: https://github.com/Islanna/DynamicReLU.

贡献:提出Dynamic ReLU激活函数

ReLU是深度神经网络中常用的激活函数。到目前为止,ReLU及其推广(非参数或参数)都是静态的,对所有的输入样本执行相同的操作。在本文中,我们提出了Dynamic ReLU激活函数(DY-ReLU),它的参数依赖于所有输入。其关键在于DY-ReLU将全局上下文编码为超函数,并相应地调整分段线性激活函数。与静态模型相比,DY-ReLU的额外计算开销可以忽略不计,但其表现能力显着提高,特别是对于轻量神经网络。仅仅通过简单地在MobileNetV2上使用DY-ReLU ,ImageNet分类的最高精度就可以从72.0%提高到76.2%,而计算量只增加了5%。

Dy-ReLU特点(优点):

  • 将所有输入元素 x={xc​} 的全局上下文编码在超参数 θ(x) 中(运用SE模块的注意力机制),以适应激活函数fθ(x)​(x)(可以根据输入数据x,动态的学习选择最佳的激活函数)。

ReLU是深度学习中很重要的里程碑,简单但强大,能够极大地提升神经网络的性能。目前也有很多ReLU的改进版,比如Leaky ReLU和 PReLU,而这些改进版和原版的最终参数都是固定的。所以论文自然而然地想到,如果能够根据输入特征来调整ReLU的参数可能会更好。

定义:

K认为是分段数

Relation to Prior Work

网络实现:

DY-ReLU的可能性很大,表1展示了DY-ReLU与原版ReLU以及其变种的关系。在学习到特定的参数后,DY-ReLU可等价于ReLU、LeakyReLU以及PReLU。而当K=1,偏置bc1=0时,则等价于SE模块。另外DY-ReLU也可以是一个动态且高效的Maxout算子,相当于将Maxout的K个卷积转换为K个动态的线性变化,然后同样地输出最大值。

DY-ReLU-A

空间位置和维度均共享(spatial and channel-shared),计算如图2a所示,仅需输出2K个参数,计算最简单,表达能力也最弱。

DY-ReLU-B

仅空间位置共享(spatial-shared and channel-wise),计算如图2b所示,输出2KC个参数。

问题:

HRNet 论文和代码详解

github: https://github.com/HRNet/HRNet-Semantic-Segmentation

Paper: https://arxiv.org/abs/1908.07919

High-Resoultion Net(HRNet)由微软亚洲研究院和中科大提出,发表在CVPR2019

摘要:高分辨率表示对于位置敏感的视觉问题十分重要,比如目标检测、语义分割、姿态估计。为了这些任务位置信息更加精准,很容易想到的做法就是维持高分辨率的feature map,事实上HRNet之前几乎所有的网络都是这么做的,通过下采样得到强语义信息,然后再上采样恢复高分辨率恢复位置信息(如下图所示),然而这种做法,会导致大量的有效信息在不断的上下采样过程中丢失。而HRNet通过并行多个分辨率的分支,加上不断进行不同分支之间的信息交互,同时达到强语义信息和精准位置信息的目的。

模型的主要特点是在整个过程中特征图(Feature Map)始终保持高分辨率,这与之前主流方法思路上有很大的不同。在HRNet之前,2D人体姿态估计算法是采用(Hourglass/CPN/Simple Baseline/MSPN等)将高分辨率特征图下采样至低分辨率,再从低分辨率特征图恢复至高分辨率的思路(单次或重复多次),以此过程实现了多尺度特征提取的一个过程。

HRNet在整个过程中保持特征图的高分辨率,但多尺度特征提取是姿态估计模型一定要实现的过程,那么HRNet是如何实现多尺度特征提取的呢?模型是通过在高分辨率特征图主网络逐渐并行加入低分辨率特征图子网络,不同网络实现多尺度融合与特征提取实现的。

特点与优势:

(1)作者提出的方法是并行连接高分辨率与低分辨率网络,而不是像之前方法那样串行连接。因此,其方法能够保持高分辨率,而不是通过一个低到高的过程恢复分辨率,因此预测的heatmap可能在空间上更精确。

(2)本文提出的模型融合相同深度和相似级别的低分辨率特征图来提高高分辨率的特征图的表示效果,并进行重复的多尺度融合。

缺点:因为特征图分辨率过大,而且数量多,这样肯定会导致巨额的耗时计算,对显存对硬件要求更高了

HRNet结构细节

Backbone设计

我将HRNet整个backbone部分进行了拆解,分成4个stage,每个stage分成蓝色框和橙色框两部分。其中蓝色框部分是每个stage的基本结构,由多个branch组成,HRNet中stage1蓝色框使用的是BottleNeck,stage2&3&4蓝色框使用的是BasicBlock。其中橙色框部分是每个stage的过渡结构,HRNet中stage1橙色框是一个TransitionLayer,stage2&3橙色框是一个FuseLayer和一个TransitionLayer的叠加,stage4橙色框是一个FuseLayer。

解释一下为什么这么设计,FuseLayer是用来进行不同分支的信息交互的,TransitionLayer是用来生成一个下采样两倍分支的输入feature map的,stage1橙色框显然没办法做FuseLayer,因为前一个stage只有一个分支,stage4橙色框后面接neck和head了,显然也不再需要TransitionLayer了。

整个backbone的构建流程可以总结为:make_backbone -> make_stages -> make_branches

有关backbone构建相关的看源码,主要讲一下FuseLayerTransitionLayerNeck的设计

FuseLayer设计

FuseLayer部分以绿色框为例,融合前为pre,融合后为post,静态构建一个二维矩阵,然后将pre和post对应连接的操作一一填入这个二维矩阵中。

以上图为例,图1的pre1和post1的操作为空,pre2和post1的操作为2倍上采,pre3和post1的操作为4倍上采;图2的pre1和post2的操作为3×3卷积下采,pre2和post2的操作为空,pre3和post2的操作为2倍上采;图3的pre1和post3的操作为连续两个3×3卷积下采,pre2和post3的操作为3×3卷积下采,pre3和post的操作为空。

前向计算时用一个二重循环将构建好的二维矩阵一一解开,将对应同一个post的pre转换后进行融合相加。比如post1 = f11(pre1) + f12(pre2) + f13(pre3)

FuseLayer的整体code如下:

def _make_fuse_layers(self):
  fuse_layers = []
  for post_index, out_channel in enumerate(self.out_channels[:len(self.in_channels)]):
      fuse_layer = []
      for pre_index, in_channel in enumerate(self.in_channels):
          if pre_index > post_index:
              fuse_layer.append(nn.Sequential(
                  nn.Conv2d(in_channel, out_channel, 1, 1, 0, bias=False),
                  nn.BatchNorm2d(out_channel, momentum=0.1),
                  nn.Upsample(scale_factor=2**(pre_index-post_index), mode='nearest')))
          elif pre_index < post_index:
              conv3x3s = []
              for cur_index in range(post_index - pre_index):
                  out_channels_conv3x3 = out_channel if cur_index == post_index - pre_index - 1 else in_channel
                  conv3x3 = nn.Sequential(
                      nn.Conv2d(in_channel, out_channels_conv3x3, 3, 2, 1, bias=False),
                      nn.BatchNorm2d(out_channels_conv3x3, momentum=0.1)
                  )
                  if cur_index < post_index - pre_index - 1:
                      conv3x3.add_module('relu_{}'.format(cur_index), nn.ReLU(False))
                  conv3x3s.append(conv3x3)
              fuse_layer.append(nn.Sequential(*conv3x3s))
          else:
              fuse_layer.append(None)
      fuse_layers.append(nn.ModuleList(fuse_layer))
  return nn.ModuleList(fuse_layers)

def forward(self, x):
  x_fuse = []
  for post_index in range(len(self.fuse_layers)):
      y = 0
      for pre_index in range(len(self.fuse_layers)):
          if post_index == pre_index:
              y += x[pre_index]
          else:
              y += self.fuse_layers[post_index][pre_index](x[pre_index])
      x_fuse.append(self.relu(y))

TransitionLayer设计

TransitionLayer以黄色框为例,静态构建一个一维矩阵,然后将pre和post对应连接的操作一一填入这个一维矩阵中。当pre1&post1、pre2&post2、pre3&post3的通道数对应相同时,一维矩阵填入None;通道数不相同时,对应位置填入一个转换卷积。post4比较特殊,这一部分代码和图例不太一致,图例是pre1&pre2&pre3都进行下采然后进行融合相加得到post4,而代码中post4通过pre3下采得到。

TransitionLayer整体code如下

def _make_transition_layers(self):
  num_branches_pre = len(self.in_channels)
  num_branches_post = len(self.out_channels)
  transition_layers = []
  for post_index in range(num_branches_post):
      if post_index < len(self.in_channels):
          if self.in_channels[post_index] != self.out_channels[post_index]:
              transition_layers.append(nn.Sequential(
                  nn.Conv2d(self.in_channels[post_index], self.out_channels[post_index], 3, 1, 1, bias=False),
                  nn.BatchNorm2d(self.out_channels[post_index], momentum=0.1),
                  nn.ReLU(inplace=True)
              ))
          else:
              transition_layers.append(None)
      else:
          conv3x3s = []
          for pre_index in range(post_index + 1 - num_branches_pre):
              in_channels_conv3x3 = self.in_channels[-1]
              out_channels_conv3x3 = self.out_channels[post_index] if pre_index == post_index - \
                  num_branches_pre else in_channels_conv3x3
              conv3x3s.append(nn.Sequential(
                  nn.Conv2d(in_channels_conv3x3, out_channels_conv3x3, 3, 2, 1, bias=False),
                  nn.BatchNorm2d(out_channels_conv3x3, momentum=0.1),
                  nn.ReLU(inplace=True)
              ))
          transition_layers.append(nn.Sequential(*conv3x3s))
  return nn.ModuleList(transition_layers)

def forward(self, x):
  x_trans = []
  for branch_index, transition_layer in enumerate(self.transition_layers):
      if branch_index < len(self.transition_layers) - 1:
          if transition_layer:
              x_trans.append(transition_layer(x[branch_index]))
          else:
              x_trans.append(x[branch_index])
      else:
          x_trans.append(transition_layer(x[-1]))

Neck设计

我把HRNet所描述的make_head过程理解成make_neck(因为一般意义上将最后的fc层理解成head更为清晰,这个在很多开源code中都是这样子拆解的)。下面着重讲解一下HRNet的neck设计。

HRNet的backbone输出有四个分支,paper中给出了几种方式对输出分支进行操作。

(a)图是HRNetV1的操作方式,只使用分辨率最高的feature map。

(b)图是HRNetV2的操作方式,将所有分辨率的feature map(小的特征图进行upsample)进行concate,主要用于语义分割和面部关键点检测。

(c)图是HRNetV2p的操作方式,在HRNetV2的基础上,使用了一个特征金字塔,主要用于目标检测。

而在图像分类任务上,HRNet有另一种特殊的neck设计

HRNet的neck可以分成三个部分,IncreLayer(橙色框),DownsampLayer(蓝色框)和FinalLayer(绿色框)。对每个backbone的输出分支进行升维操作,然后按照分辨率从大到小依次进行下采样同时从上到下逐级融合相加,最后用一个1x1conv升维。

def _make_neck(self, in_channels):
  head_block = Bottleneck
  self.incre_channels = [32, 64, 128, 256]
  self.neck_out_channels = 2048

  incre_modules = []
  downsamp_modules = []
  num_branches = len(self.in_channels)
  for index in range(num_branches):
      incre_module = self._make_layer(head_block, in_channels[index], incre_channels[index], 1, stride=1)
      incre_modules.append(incre_module)
      if index < num_branches - 1:
          downsamp_in_channels = self.incre_channels[index] * incre_module.expansion
          downsamp_out_channels = self.incre_channels[index+1] * incre_module.expansion
          downsamp_module = nn.Sequential(
              nn.Conv2d(in_channels=downsamp_in_channels, out_channels=downsamp_out_channels,
                        kernel_size=3, stride=2, padding=1),
              nn.BatchNorm2d(downsamp_out_channels, momentum=0.1),
              nn.ReLU(inplace=True)
          )
          downsamp_modules.append(downsamp_module)
  incre_modules = nn.ModuleList(incre_modules)
  downsamp_modules = nn.ModuleList(downsamp_modules)
  final_layer = nn.Sequential(
      nn.Conv2d(in_channels=self.out_channels[-1] * 4, out_channels=2048,
                kernel_size=1, stride=1, padding=0),
      nn.BatchNorm2d(2048, momentum=0.1),
      nn.ReLU(inplace=True)
  )
  return incre_modules, downsamp_modules, fine_layer

def forward(self, x):
  y = self.incre_modules[0](x[0])
  for index in range(len(self.downsamp_modules)):
      y = self.incre_modules[index+1](x[index+1]) + self.downsamp_modules[index](y)
  y = self.final_layer(y)
  y = F.avg_pool2d(y, kernel_size=y.size()[2:]).view(y.size(0), -1)

还有几个小细节

  1. BN层的momentom都设置为0.1
  2. stem使用的是两层stried为2的conv3x3
  3. FuseLayer的ReLU的inplace都设置为False

Swin Transformer论文解读与思考

论文: https://arxiv.org/abs/2103.14030

github:https://github.com/microsoft/Swin-Transformer

代码详解 https://zhuanlan.zhihu.com/p/384514268

Vision Transformer , Vision MLP 超详细解读 (原理分析+代码解读) (目录)

论文详解:https://space.bilibili.com/1567748478/channel/collectiondetail?sid=32744

Swin Transformer视频讲解:

https://github.com/WZMIAOMIAO/deep-learning-for-image-processing

摘要 :目前transformer应用于CV领域的挑战主要有两个,一个是图片多尺度语义信息的问题,同一个物体在不同图片中的大小尺度变化很大,另外就是难以处理高分辨的图片,如果以pix像素作为序列元素,那么计算成本太大,因此一部分方法是将CNN提取图片特征在送进transformer中 ,或者通过patch,将图片变成一个个的patch。作者提出Swin Transformer 目标是希望作为一种计算机视觉的通用主干网络(因为VIT的提出已经证明了Transformer在CV的可行性),这是一种层级的架构。通过窗口注意力以及转移窗口注意力,不仅降低了计算量,同时层级架构对于不同尺度的信息处理都十分灵活,该架构在图像分类、目标检测、语义分割等任务中表现出色。(对于图像分类、目标检测、语义分割等下游任务,尤其是密集预测任务,多尺度特征是十分必要的)

引言

首先来看作者给出的Swin Transformer 和 VIT结构对比:

VIT的patch固定16*16(可以认为是16倍下采样),多尺度特征处理不好,因为整个过程都是在同一尺度下操作的,出来的特征是单尺度的,优点是全局的特征处理比较强,因为是在全局的尺度进行操作的,但因此他的复杂度跟图像尺寸成平方倍的增长,很难处理目前图像分割检测。再来看 Swin Transformer ,通过图可以看出作者借鉴了CNN的很多设计思路,为了减少序列长度,减低计算量,仅在上面的红框中进行自注意力计算,计算复杂度会跟整张图片的大小成线性关系。另外作者使用基于窗口的注意力的也可以很好的把握物体的全局信息(因为在CV中,一个物体的绝大部分都存在单个windows窗口中,很少会横跨多个窗口),另外CNN网络的如何抓住物体的多尺度特征?是因为pool池化层的存在,每次池化能够增大卷积核看到的感受野。因此作者提出了patch merging,将相邻的四个patch合并成一个大patch(可以认为是加权池化),这样合并出来的一个 大patch就可以看到四个小patch内容感受野增大。有了多尺度特征(4*,8*,16*多尺度特征图)以后,可以接一个FPN头,由于做检测任务,也可以放在unet做分割任务,这就是作者所说的, Swin Transformer 是可以做一个通用骨干网络。

    Transformer的初衷就是更好的理解上下文,如果窗口都是不重叠的,那自注意力真的就变成孤立自注意力,就没有全局建模的能力   
    Swin Transformer 的一个关键设计因素:移动窗口操作。在第l层,通过划分不同的小窗口(实际中是一个窗口有7*7个patch(最小单位),这里示意图以4*4的patch作为一个窗口),自注意力只在窗口中计算 ,就可以有效降低序列长度,从而减少计算复杂度。shift操作可以认为是将l层的窗口整体向右下加移动两个patch所形成的新的窗口,新的特征图进行分割windws以后就有l+1层所示的这些窗口(如下图共九个)了。如果没有shift,那么所有窗口不重叠,在窗口进行自注意力时候,窗口之间无法交互,就无法达到transformer的初衷了(更好的理解上下文),shift后不同窗口的patch就可以进行交互了。再加上一个patch merging操作,不断扩大感受野,到最后几层的时候,每个patch的感受野已经很大了,实际上就可以看到大部分图片了,shift操作以后,就可以看成是全局注意力操作,这样即省内存效果也好。

引言的最后,作者坚信,一个CV和NLP大一统的框架是可以促进两个领域共同发展的,但实际上 Swin Transformer 更多的是利用了CNN的先验知识,从而在计算机视觉领域大杀四方。但是在模型大一统上,也就是 unified architecture 上来说,其实 ViT 还是做的更好的,因为它真的可以什么都不改,什么先验信息都不加,就能让Transformer在两个领域都能用的很好,这样模型不仅可以共享参数,而且甚至可以把所有模态的输入直接就拼接起来,当成一个很长的输入,直接扔给Transformer去做,而不用考虑每个模态的特性

先看结论:

这篇论文提出了 Swin Transformer,它是一个层级式的Transformer,而且它的计算复杂度是跟输入图像的大小呈线性增长的。Swin Transformerr 在 COCO 和 ADE20K上的效果都非常的好,远远超越了之前最好的方法,所以作者说基于此,希望 Swin Transformer 能够激发出更多更好的工作,尤其是在多模态方面。

因为在Swin Transformer 这篇论文里最关键的一个贡献就是基于 Shifted Window 的自注意力,它对很多视觉的任务,尤其是对下游密集预测型的任务是非常有帮助的,但是如果 Shifted Window 操作不能用到 NLP 领域里,其实在模型大一统上论据就不是那么强了,所以作者说接下来他们的未来工作就是要把 Shifted Windows用到 NLP 里面,而且如果真的能做到这一点,那 Swin Transformer真的就是一个里程碑式的工作了,而且模型大一统的故事也就讲的圆满了

方法

主要分为两大块

  • 大概把整体的流程讲了一下,主要就是过了一下前向过程,以及提出的 patch merging 操作是怎么做的
  • 基于 Shifted Window 的自注意力,Swin Transformer怎么把它变成一个transformer block 进行计算

前向过程

  • 假设说有一张224*224*3(ImageNet 标准尺寸)的输入图片
  • 第一步就是像 ViT 那样把图片打成 patch,在 Swin Transformer 这篇论文里,它的 patch size 是4*4,而不是像 ViT 一样16*16,所以说它经过 patch partition 打成 patch 之后,得到图片的尺寸是56*56*48,56就是224/4,因为 patch size 是4,向量的维度48,因为4*4*3,3 是图片的 RGB 通道
  • 打完了 patch ,接下来就要做 Linear Embedding,也就是说要把向量的维度变成一个预先设置好的值,就是 Transformer 能够接受的值,在 Swin Transformer 的论文里把这个超参数设为 c,对于 Swin tiny 网络来说,也就是上图中画的网络总览图,它的 c 是96,所以经历完 Linear Embedding 之后,输入的尺寸就变成了56*56*96,前面的56*56就会拉直变成3136,变成了序列长度,后面的96就变成了每一个token向量的维度,其实 Patch Partition 和 Linear Embedding 就相当于是 ViT 里的Patch Projection 操作,而在代码里也是用一次卷积操作就完成了,
  • 第一部分跟 ViT 其实还是没有区别的,但紧接着区别就来了
  • 首先序列长度是3136,对于 ViT 来说,用 patch size 16*16,它的序列长度就只有196,是相对短很多的,这里的3136就太长了,是目前来说Transformer不能接受的序列长度,所以 Swin Transformer 就引入了基于窗口的自注意力计算,每个窗口按照默认来说,都只有七七四十九个 patch,所以说序列长度就只有49就相当小了,这样就解决了计算复杂度的问题
  • 所以也就是说, stage1中的swin transformer block 是基于窗口计算自注意力的,现在暂时先把 transformer block当成是一个黑盒,只关注输入和输出的维度,对于 Transformer 来说,如果不对它做更多约束的话,Transformer输入的序列长度是多少,输出的序列长度也是多少,它的输入输出的尺寸是不变的,所以说在 stage1 中经过两层Swin Transformer block 之后,输出还是56*56*96
  • 到这其实 Swin Transformer的第一个阶段就走完了,也就是先过一个 Patch Projection 层,然后再过一些 Swin Transformer block,接下来如果想要有多尺寸的特征信息,就要构建一个层级式的 transformer,也就是说需要一个像卷积神经网络里一样,有一个类似于池化的操作

Patch Merging

Patch Merging 其实在之前一些工作里也有用到,它很像 Pixel Shuffle 的上采样的一个反过程,Pixel Shuffle 是 lower level 任务中很常用的一个上采样方式

  • 假如有一个张量, Patch Merging 顾名思义就是把临近的小 patch 合并成一个大 patch,这样就可以起到下采样一个特征图的效果了
  • 这里因为是想下采样两倍,所以说在选点的时候是每隔一个点选一个,也就意味着说对于这个张量来说,每次选的点是1、1、1、1
  • 其实在这里的1、2、3、4并不是矩阵里有的值,而是给它的一个序号,同样序号位置上的 patch 就会被 merge 到一起,这个序号只是为了帮助理解
  • 经过隔一个点采一个样之后,原来的这个张量就变成了四个张量,也就是说所有的1都在一起了,2在一起,3在一起,4在一起,如果原张量的维度是 h * w * c ,当然这里 c 没有画出来,经过这次采样之后就得到了4个张量,每个张量的大小是 h/2、w/2,它的尺寸都缩小了一倍
  • 现在把这四个张量在 c 的维度上拼接起来,也就变成了下图中所画出来的形式,张量的大小就变成了 h/2 * w/2 * 4c,相当于用空间上的维度换了更多的通道数
  • 通过这个操作,就把原来一个大的张量变小了,就像卷积神经网络里的池化操作一样,为了跟卷积神经网络那边保持一致(不论是 VGGNet 还是 ResNet,一般在池化操作降维之后,通道数都会翻倍,从128变成256,从256再变成512),所以这里也只想让他翻倍,而不是变成4倍,所以紧接着又再做了一次操作,就是在 c 的维度上用一个1乘1的卷积,把通道数降下来变成2c,通过这个操作就能把原来一个大小为 h*w*c 的张量变成 h/2 * w/2 *2c 的一个张量,也就是说空间大小减半,但是通道数乘2,这样就跟卷积神经网络完全对等起来了

这里其实会发现,特征图的维度真的跟卷积神经网络好像,因为如果回想残差网络的多尺寸的特征,就是经过每个残差阶段之后的特征图大小也是56*56、28*28、14*14,最后是7*7

而且为了和卷积神经网络保持一致,Swin Transformer这篇论文并没有像 ViT 一样使用 CLS token,ViT 是给刚开始的输入序列又加了一个 CLS token,所以这个长度就从196变成了197,最后拿 CLS token 的特征直接去做分类,但 Swin Transformer 没有用这个 token,它是像卷积神经网络一样,在得到最后的特征图之后用global average polling,就是全局池化的操作,直接把7*7就取平均拉直变成1了

作者这个图里并没有画,因为 Swin Transformer的本意并不是只做分类,它还会去做检测和分割,所以说它只画了骨干网络的部分,没有去画最后的分类头或者检测头,但是如果是做分类的话,最后就变成了1*768,然后又变成了1*1,000

所以看完整个前向过程之后,就会发现 Swin Transformer 有四个 stage,还有类似于池化的 patch merging 操作,自注意力还是在小窗口之内做的以及最后还用的是 global average polling,所以说 Swin Transformer 这篇论文真的是把卷积神经网络和 Transformer 这两系列的工作完美的结合到了一起,也可以说它是披着Transformer皮的卷积神经网络

主要贡献

这篇论文的主要贡献就是基于窗口或者移动窗口的自注意力,这里作者又写了一段研究动机,就是为什么要引入窗口的自注意力,其实跟之前引言里说的都是一个事情,就是说全局自注意力的计算会导致平方倍的复杂度,同样当去做视觉里的下游任务,尤其是密集预测型的任务,或者说遇到非常大尺寸的图片时候,这种全局算自注意力的计算复杂度就非常贵了,所以就用窗口的方式去做自注意力

重点:窗口注意力

原图片会被平均的分成一些没有重叠的窗口,拿第一层之前的输入来举例,它的尺寸就是56*56*96,也就说有一个维度是56*56张量,然后把它切成一些不重叠的方格(论文中使用7*7的patch作为一个window窗口)

  • 现在所有自注意力的计算都是在这些小窗口里完成的,就是说序列长度永远都是7*7=49
  • 原来大的整体特征图到底里面会有多少个窗口呢?其实也就是每条边56/7就8个窗口,也就是说一共会有8*8等于64个窗口,就是说会在这64个窗口里分别去算它们的自注意力

基于窗口的自注意力模式的计算复杂度计算:

  • 如果现在有一个输入,自注意力首先把它变成 q k v 三个向量,这个过程其实就是原来的向量分别乘了三个系数矩阵
  • 一旦得到 query 和 k 之后,它们就会相乘,最后得到 attention,也就是自注意力的矩阵
  • 有了自注意力之后,就会和 value 做一次乘法,也就相当于是做了一次加权
  • 最后因为是多头自注意力,所以最后还会有一个 projection layer,这个投射层会把向量的维度投射到我们想要的维度

如果这些向量都加上它们该有的维度,也就是说刚开始输入是 h*w*c

  • 公式(1)对应的是标准的多头自注意力的计算复杂度
  • 每一个图片大概会有 h*w 个 patch,在刚才的例子里,h 和 w 分别都是56,c 是特征的维度
  • 公式(2)对应的是基于窗口的自注意力计算的复杂度,这里的 M 就是刚才的7,也就是说一个窗口的某条边上有多少个patch

基于窗口的自注意力计算复杂度又是如何得到的呢?

  • 因为在每个窗口里算的还是多头自注意力,所以可以直接套用公式(1),只不过高度和宽度变化了,现在高度和宽度不再是 h * w,而是变成窗口有多大了,也就是 M*M,也就是说现在 h 变成了 M,w 也是 M,它的序列长度只有 M * M 这么大
  • 所以当把 M 值带入到公式(1)之后,就得到计算复杂度是4 * M^2 * c^2 + 2 * M^4 * c,这个就是在一个窗口里算多头自注意力所需要的计算复杂度
  • 那我们现在一共有 h/M * w/M 个窗口,现在用这么多个窗口乘以每个窗口所需要的计算复杂度就能得到公式(2)了

对比公式(1)和公式(2),虽然这两个公式前面这两项是一样的,只有后面从 (h*w)^2变成了 M^2 * h * w,看起来好像差别不大,但其实如果仔细带入数字进去计算就会发现,计算复杂的差距是相当巨大的,因为这里的 h*w 如果是56*56的话, M^2 其实只有49,所以是相差了几十甚至上百倍的

这种基于窗口计算自注意力的方式虽然很好地解决了内存和计算量的问题,但是窗口和窗口之间没有通信,这样就达不到全局建模了,也就文章里说的会限制模型的能力,所以最好还是要有一种方式能让窗口和窗口之间互相通信起来,这样效果应该会更好,因为具有上下文的信息,所以作者就提出移动窗口的方式

移动窗口:

移动窗口就是把原来的窗口往右下角移动一半窗口的距离,如果Transformer是上下两层连着做这种操作,先是 window再是 shifted window 的话,就能起到窗口和窗口之间互相通信的目的了

所以说在 Swin Transformer里, transformer block 的安排是有讲究的,每次都是先要做一次基于窗口的多头自注意力,然后再做一次基于移动窗口的多头自注意力,这样就达到了窗口和窗口之间的互相通信。如下图所示

  • 每次输入先进来之后先做一次 Layernorm,然后做窗口的多头自注意力,然后再过 Layernorm 过 MLP,第一个 block 就结束了
  • 这个 block 结束以后,紧接着做一次Shifted window,也就是基于移动窗口的多头自注意力,然后再过 MLP 得到输出
  • 这两个 block 加起来其实才算是 Swin Transformer 一个基本的计算单元,这也就是为什么stage1、2、3、4中的 swin transformer block 为什么是 *2、*2、*6、*2,也就是一共有多少层 Swin Transformer block 的数字总是偶数,因为它始终都需要两层 block连在一起作为一个基本单元,所以一定是2的倍数

到此,Swin Transformer整体的故事和结构就已经讲完了,主要的研究动机就是想要有一个层级式的 Transformer,为了这个层级式,所以介绍了 Patch Merging 的操作,从而能像卷积神经网络一样把 Transformer 分成几个阶段,为了减少计算复杂度,争取能做视觉里密集预测的任务,所以又提出了基于窗口和移动窗口的自注意力方式,也就是连在一起的两个Transformer block,最后把这些部分加在一起,就是 Swin Transformer 的结构

提高移动窗口的计算效率:

  • 一个是怎样提高移动窗口的计算效率,他们采取了一种非常巧妙的 masking(掩码)的方式
  • 另外一个点就是这篇论文里没有用绝对的位置编码,而是用相对的位置编码

masking(掩码)的方式计算移动窗口自注意力:为什么需要使用?

为了提高计算效率,因为如果直接计算右下图的九个窗口的自注意力,不同大小的窗口无法合并成一个batch进行计算。

  • 上图是一个基础版本的移动窗口,就是把左边的窗口模式变成了右边的窗口方式
  • 虽然这种方式已经能够达到窗口和窗口之间的互相通信了,但是会发现一个问题,就是原来计算的时候,特征图上只有四个窗口,但是做完移动窗口操作之后得到了9个窗口,窗口的数量增加了,而且每个窗口里的元素大小不一,比如说中间的窗口还是4*4,有16个 patch,但是别的窗口有的有4个 patch,有的有8个 patch,都不一样了,如果想做快速运算,就是把这些窗口全都压成一个 patch直接去算自注意力,就做不到了,因为窗口的大小不一样
  • 有一个简单粗暴的解决方式就是把这些小窗口周围再 pad 上0 ,把它照样pad成和中间窗口一样大的窗口,这样就有9个完全一样大的窗口,这样就还能把它们压成一个batch,就会快很多
  • 但是这样的话,无形之中计算复杂度就提升了,因为原来如果算基于窗口的自注意力只用算4个窗口,但是现在需要去算9个窗口,复杂度一下提升了两倍多,所以还是相当可观的
  • 那怎么能让第二次移位完的窗口数量还是保持4个,而且每个窗口里的patch数量也还保持一致呢?作者提出了一个非常巧妙的掩码方式,如下图所示

上图是说,当通过普通的移动窗口方式,得到9个窗口之后,现在不在这9个窗口上算自注意力,先再做一次循环移位( cyclic shift )

  • 经过这次循环移位之后,原来的窗口(虚线)就变成了现在窗口(实线)的样子,那如果在大的特征图上再把它分成四宫格的话,我在就又得到了四个窗口,意思就是说移位之前的窗口数也是4个,移完位之后再做一次循环移位得到窗口数还是4个,这样窗口的数量就固定了,也就说计算复杂度就固定了
  • 但是新的问题就来了,虽然对于移位后左上角的窗口(也就是移位前最中间的窗口)来说,里面的元素都是互相紧挨着的,他们之间可以互相两两做自注意力,但是对于剩下几个窗口来说,它们里面的元素是从别的很远的地方搬过来的,所以他们之间,按道理来说是不应该去做自注意力,也就是说他们之间不应该有什么太大的联系
  • 解决这个问题就需要一个很常规的操作,也就是掩码操作,这在Transformer过去的工作里是层出不穷,很多工作里都有各式各样的掩码操作
  • 在 Swin Transformer这篇论文里,作者也巧妙的设计了几种掩码的方式,从而能让一个窗口之中不同的区域之间也能用一次前向过程,就能把自注意力算出来,但是互相之间都不干扰,也就是后面的 masked Multi-head Self Attention(MSA)
  • 算完了多头自注意力之后,还有最后一步就是需要把循环位移再还原回去,也就是说需要把A、B、C再还原到原来的位置上去,原因是还需要保持原来图片的相对位置大概是不变的,整体图片的语义信息也是不变的,如果不把循环位移还原的话,那相当于在做Transformer的操作之中,一直在把图片往右下角移,不停的往右下角移,这样图片的语义信息很有可能就被破坏掉了
  • 所以说整体而言,上图介绍了一种高效的、批次的计算方式比如说本来移动窗口之后得到了9个窗口,而且窗口之间的patch数量每个都不一样,为了达到高效性,为了能够进行批次处理,先进行一次循环位移,把9个窗口变成4个窗口,然后用巧妙的掩码方式让每个窗口之间能够合理地计算自注意力,最后再把算好的自注意力还原,就完成了基于移动窗口的自注意力计算

掩码操作如何实现 :

作者通过这种巧妙的循环位移的方式和巧妙设计的掩码模板,从而实现了只需要一次前向过程,就能把所有需要的自注意力值都算出来,而且只需要计算4个窗口,也就是说窗口的数量没有增加,计算复杂度也没有增加,非常高效的完成了这个任务

作者给出了不同窗口的不同掩码矩阵:

上图示例的Cyclic Shifting方法,可以保持面向计算的window数量保持不变(还是2X2),在window内部通过attention mask来计算子window中的自注意力。

Swin Transformer的几个变体

  • Swin Tiny
  • Swin Small
  • Swin Base
  • Swin Large

Swin Tiny的计算复杂度跟 ResNet-50 差不多,Swin Small 的复杂度跟 ResNet-101 是差不多的,这样主要是想去做一个比较公平的对比

这些变体之间有哪些不一样呢?,其实主要不一样的就是两个超参数

  • 一个是向量维度的大小 c
  • 另一个是每个 stage 里到底有多少个 transform block

这里其实就跟残差网络就非常像了,残差网络也是分成了四个 stage,每个 stage 有不同数量的残差块

实验

分类

首先是分类上的实验,这里一共说了两种预训练的方式

  • 第一种就是在正规的ImageNet-1K(128万张图片、1000个类)上做预训练
  • 第二种方式是在更大的ImageNet-22K(1,400万张图片、2万多个类别)上做预训练

当然不论是用ImageNet-1K去做预训练,还是用ImageNet-22K去做预训练,最后测试的结果都是在ImageNet-1K的测试集上去做的,结果如下表所示

  • 上半部分是ImageNet-1K预训练的模型结果
  • 下半部分是先用ImageNet-22K去预训练,然后又在ImageNet-1K上做微调,最后得到的结果
  • 在表格的上半部分,作者先是跟之前最好的卷积神经网络做了一下对比,RegNet 是之前 facebook 用 NASA 搜出来的模型,EfficientNet 是 google 用NASA 搜出来的模型,这两个都算之前表现非常好的模型了,他们的性能最高会到 84.3
  • 接下来作者就写了一下之前的 Vision Transformer 会达到什么效果,对于 ViT 来说,因为它没有用很好的数据增强,而且缺少偏置归纳,所以说它的结果是比较差的,只有70多
  • 换上 DeiT 之后,因为用了更好的数据增强和模型蒸馏,所以说 DeiT Base 模型也能取得相当不错的结果,能到83.1
  • 当然 Swin Transformer 能更高一些,Swin Base 最高能到84.5,稍微比之前最好的卷积神经网络高那么一点点,就比84.3高了0.2
  • 虽然之前表现最好的 EfficientNet 的模型是在 600*600 的图片上做的,而 Swin Base 是在 384*384 的图片上做的,所以说 EfficientNet 有一些优势,但是从模型的参数和计算的 FLOPs 上来说 EfficientNet 只有66M,而且只用了 37G 的 FLOPs,但是 Swin Transformer 用了 88M 的模型参数,而且用了 47G 的 FLOPs,所以总体而言是伯仲之间
  • 表格的下半部分是用 ImageNet-22k 去做预训练,然后再在ImageNet-1k上微调最后得到的结果
  • 这里可以看到,一旦使用了更大规模的数据集,原始标准的 ViT 的性能也就已经上来了,对于 ViT large 来说它已经能得到 85.2 的准确度了,已经相当高了
  • 但是 Swin Large 更高,Swin Large 最后能到87.3,这个是在不使用JFT-300M,就是特别大规模数据集上得到的结果,所以还是相当高的

目标检测

  • 表2(a)中测试了在不同的算法框架下,Swin Transformer 到底比卷积神经网络要好多少,主要是想证明 Swin Transformer 是可以当做一个通用的骨干网络来使用的,所以用了 Mask R-CNN、ATSS、RepPointsV2 和SparseR-CNN,这些都是表现非常好的一些算法,在这些算法里,过去的骨干网络选用的都是 ResNet-50,现在替换成了 Swin Tiny
  • Swin Tiny 的参数量和 FLOPs 跟 ResNet-50 是比较一致的,从后面的对比里也可以看出来,所以他们之间的比较是相对比较公平的
  • 可以看到,Swin Tiny 对 ResNet-50 是全方位的碾压,在四个算法上都超过了它,而且超过的幅度也是比较大的
  • 接下来作者又换了一个方式做测试,现在是选定一个算法,选定了Cascade Mask R-CNN 这个算法,然后换更多的不同的骨干网络,比如 DeiT-S、ResNet-50 和 ResNet-101,也分了几组,结果如上图中表2(b)所示
  • 可以看出,在相似的模型参数和相似的 Flops 之下,Swin Transformer 都是比之前的骨干网络要表现好的
  • 接下来作者又做了第三种测试的方式,如上图中的表2(c)所示,就是系统层面的比较,这个层面的比较就比较狂野了,就是现在追求的不是公平比较,什么方法都可以上,可以使用更多的数据,可以使用更多的数据增强,甚至可以在测试的使用 test time augmentation(TTA)的方式
  • 可以看到,之前最好的方法 Copy-paste 在 COCO Validation Set上的结果是55.9,在 Test Set 上的结果是56,而这里如果跟最大的 Swin Transformer–Swin Large 比,它的结果分别能达到58和58.7,这都比之前高了两到三个点

语义分割

  • 上图表3里可以看到之前的方法,一直到 DeepLab V3、ResNet 其实都用的是卷积神经网络,之前的这些方法其实都在44、45左右徘徊
  • 但是紧接着 Vision Transformer 就来了,那首先就是 SETR 这篇论文,他们用了 ViT Large,所以就取得了50.3的这个结果
  • Swin Transformer Large也取得了53.5的结果,就刷的更高了
  • 其实作者这里也有标注,就是有两个“+”号的,意思是说这些模型是在ImageNet-22K 数据集上做预训练,所以结果才这么好

消融实验

实验结果如下图所示

  • 上图中表4主要就是想说一下移动窗口以及相对位置编码到底对 Swin Transformer 有多有用
  • 可以看到,如果光分类任务的话,其实不论是移动窗口,还是相对位置编码,它的提升相对于基线来说,也没有特别明显,当然在ImageNet的这个数据集上提升一个点也算是很显着了
  • 但是他们更大的帮助,主要是出现在下游任务里,就是 COCO 和 ADE20K 这两个数据集上,也就是目标检测和语义分割这两个任务上
  • 可以看到,用了移动窗口和相对位置编码以后,都会比之前大概高了3个点左右,提升是非常显着的,这也是合理的,因为如果现在去做这种密集型预测任务的话,就需要特征对位置信息更敏感,而且更需要周围的上下文关系,所以说通过移动窗口提供的窗口和窗口之间的互相通信,以及在每个 Transformer block都做更准确的相对位置编码,肯定是会对这类型的下游任务大有帮助的

总结

虽然前面已经说了很多 Swin Transformer 的影响力啊已经这么巨大了,但其实他的影响力远远不止于此,论文里这种对卷积神经网络,对 Transformer,还有对 MLP 这几种架构深入的理解和分析是可以给更多的研究者带来思考的,从而不仅可以在视觉领域里激发出更好的工作,而且在多模态领域里,相信它也能激发出更多更好的工作

BPR:用于实例分割的边界Patch优化(CVPR2021)

 

Look Closer to Segment Better: Boundary Patch Refinement for Instance Segmentation

代码链接:https://github.com/tinyalpha/BPR

后处理分割结果,效果是即插即用后处理模块当年的sota通过将 BPR 框架应用于 PolyTransform + SegFix 基线,我们在 Cityscapes 排行榜上排名第一。

从目前的排名来说(22.09.23),排名第五,与top1相差不到2个百分点,而 BPR后处理使得PolyTransform + SegFix的效果提升了1.5个百分点。 相比于MASK-RCNN提升了4.2个百分点。

CVPR21上一篇关于实例分割的文章。对于Mask RCNN来说,其最终得到的mask分辨率太低,因此还原到原尺寸的时候,一些boundary信息就显得非常粗糙,导致预测生成的mask效果不尽如人意。而且处于boundary的pixel本身数量相比于整张image来说很少,同时本身难以做分类。现有的一些方法试图提升boundary quality,但预测mask边界这个task本身的复杂度和segmentation很接近了,因此开销较大。

因此本文作者提出了一种crop-and-refine的策略。首先通过经典的实例分割网络(如Mask RCNN)得到coarse mask。随后在mask的boundary出提取出一系列的patch,随后将这些patch送入一个Refinement Network,这个Refinement Network负责做二分类的语义分割,进而对boundary处的patch进行优化,整个后处理的优化网络称为BPR(Boundary Patch Refinement)。该网络可以解决传统Mask RCNN预测的mask的边界粗糙的问题。

本文的核心就是在Mask RCNN一类的网络给出coarse mask后,如何设计Refine Network来对这个粗糙 mask 的边界进行优化,进而得到resolution更高,boundary quality更好的mask。

给定一个coarse mask(上图a),首先需要决定这个mask的哪些部分要做refine。这里作者提出了一种sliding-window式的方法提取到boundary处的一系列patch(上图b)。具体来说,就是在mask边界处密集assign正方形的bounding box,这些box内部囊括了boundary pixel。随后,由于这些box有的overlap太大导致redundant(冗余),这里采用NMS进行过滤(上图c),以实现速度和精度的trade-off(平衡)。

随后这些survive下来的image patch(上图d)和mask patch(上图e)都resize到同一尺寸,一起喂入Refinement Network。这里作者argue说一定要喂入mask patch,因为一旦拥有mask patch的location和semantic信息,这个refinement network就不再需要学习instance-level semantic(实例类别信息,比如这个image patch属于哪个类别)了。所以,refinement network只需要学习boundary处的hard pixel,并把它们正确分类。

关于Refinement Network,其任务是为每一个提取出来的boundary patch独立地做二分类语义分割,任何的语义分割模型都可以搬过来做这个task。输入的通道数为4(RGB+mask),输出通道数为2(BG or FG),这里作者采用了HRNetV2(CVPR 2019),这种各种level feature不断做融合的网络可以maintain高分辨率的representation。通过合理的增加input size,boundary batch就可以得到比之前方法更高的resolution。

HRNetV2 网络结构

在对每个patch独立地refine以后,需要将它们reassemble(组装)到coarse mask上面。有的相邻的patch可能存在overlap的情况,最终的结果是取平均,以0.5作为阈值判断某个pixel属于前景或是背景。

Experiment

这里的指标是AP (Average precision):指的是PR曲线的面积(AP就是平均精准度,简单来说就是对PR曲线上的Precision值求均值。)对于实例分割的评价指标:使用AP评价指标

实例分割和目标检测mAP计算时除了IOU计算方式(实例分割是mask间的IOU)不同,其他都是一样的.

对于一个二分类任务,二分类器的预测结果可分为以下4类:

二分类器的结果可分为4类

Precision的定义为:

Recall的定义为: 

Precision从预测结果角度出发,描述了二分类器预测出来的正例结果中有多少是真实正例,即该二分类器预测的正例有多少是准确的;Recall从真实结果角度出发,描述了测试集中的真实正例有多少被二分类器挑选了出来,即真实的正例有多少被该二分类器召回。

逐步降低二分类器预测正例的门槛,则每次可以计算得到当前的Precision和Recall。以Recall作为横轴,Precision作为纵轴可以得到Precision-Recall曲线图,简称为P-R图。

详细解释:目标检测/实例分割中 AP 和 mAP 的混淆指标

preview

首先通过实验证明了将mask patch一并作为输入的重要性:

patch size、不同的patch extraction策略,input size对结果的影响:

RefineNet的选取,NMS的阈值:

Cityscape上与其他方法的比较:PolyTransform + SegFix baseline,达到最高的AP。

迁移到其他model上面的结果 and coco数据集上的结果

Mask-RCNN论文

论文:http://cn.arxiv.org/pdf/1703.06870v3

代码:https://github.com/facebookresearch/maskrcnn-benchmark

B站网络详解 FPN

Introduction

我们提出了一个简单、灵活、通用的实例分割框架,称为Mask R-CNN。我们的方法能够有效检测图像中的目标,同时为每个实例生成高质量的分割掩码。Mask R-CNN通过添加一个预测对象掩码的分支,与现有的边框识别分支并行,扩展了之前的Faster R-CNN。Mask R-CNN的训练很简单,只为Faster R-CNN增加了一小部分开销,运行速度为5帧/秒。此外,Mask R-CNN很容易泛化到其他任务,如人体姿态估计。我们展示了Mask R-CNN在COCO挑战赛的实例分割、目标检测和人物关键点检测任务上的最优结果。在不使用花哨技巧的情况下,Mask R-CNN在各项任务上都优于现有的单一模型,包括COCO 2016挑战赛的冠军。我们希望Mask R-CNN能够成为一个坚实的基线,并有助于简化未来实例识别的研究。

Fast/Faster R-CNN和Fully Convolutional Network(FCN)框架极大地推动了计算机视觉领域中目标检测和语义分割等方向的发展。这些方法的概念很直观,具有良好的灵活性和鲁棒性,并且能够快速训练和推理。我们这项工作的目标是为实例分割任务开发一个相对可行的框架。

实例分割具有一定的挑战性,因为它需要正确检测图像中的所有对象,同时还要精确分割每个实例。因此,它结合了目标检测和语义分割等计算机视觉任务中的元素。目标检测旨在对单个物体进行分类,并使用边框对每个物体进行定位。语义分割旨在将每个像素归类到一组固定的类别,而不区分对象实例。鉴于此,人们可能会认为需要一套复杂的方法才能获得良好的结果。然而,我们证明了一个令人惊讶的事实:简单、灵活、快速的系统也可以超越现有的最先进的实例分割模型。

我们的方法称为Mask R-CNN,通过在每个RoI(感兴趣区域,Region of Interest)上添加一个预测分割掩码的分支来扩展Faster R-CNN,并与现有的用于分类和边框回归的分支并行。掩码分支是应用于每个RoI的一个小FCN,以像素到像素的方式预测分割掩码,并且只会增加较小的计算开销。Mask R-CNN是基于Faster R-CNN框架而来的,易于实现和训练,有助于广泛、灵活的架构设计。

原则上,Mask R-CNN是Faster R-CNN的直观扩展,但正确构建掩码分支对于获得好的结果至关重要。最重要的是,Faster R-CNN的设计没有考虑网络输入和输出之间的像素到像素的对齐。这一点在RoIPool(处理实例的核心操作)如何执行粗空间量化来提取特征上表现得最为明显。为了修正错位,我们提出了一个简单的、没有量化的层,称为RoIAlign,它忠实地保留了精确的空间位置。尽管这看起来是一个很小的变化,但是RoIAlign有很大的影响:它将掩码精度提高了10%-50%,在更严格的localization指标下显示出更大的收益。其次,我们发现有必要将掩码和类别预测解耦:我们为每个类别独立预测一个二进制掩码,类别之间没有竞争,并依靠网络的RoI分类分支来预测类别。相比之下,FCN通常执行逐像素的多分类操作,将分割和分类耦合在一起,我们的实验结果表明这种方法的实例分割效果不佳。

在不使用花哨技巧的情况下,Mask R-CNN在COCO实例分割任务上就超越了之前的所有SOTA单模型,包括COCO 2016比赛的冠军。作为副产品,我们的方法在COCO目标检测任务上也表现出色。在消融实验中,我们评估了多个基本实例,这使我们能够证明Mask R-CNN的鲁棒性,并分析其核心因素的影响。

我们的模型可以在GPU上以每帧约200ms的速度运行,在一台8-GPU的机器上进行COCO训练需要1-2天。我们相信,快速的训练和测试,以及框架的灵活性和准确性,将有利于未来实例分割的研究。

最后,我们通过COCO关键点数据集上的人体姿态估计任务展示了Mask R-CNN框架的通用性。通过将每个关键点视为一个独热二进制掩码,只需对Mask R-CNN稍加修改,即可用于检测特定实例的姿态。Mask R-CNN超越了COCO 2016关键点检测比赛的冠军,并且能够以5帧/秒的速度运行。因此,Mask R-CNN可以被更广泛地视为一个实例识别的灵活框架,并且很容易泛化到其他更复杂的任务上。

模型方法

Mask R-CNN方法很简单:Faster R-CNN对每个候选对象有两个输出,一个是类别标签,另一个是边框偏移量。在此基础上,我们添加了第三个分支,用于输出分割掩码。因此,Mask R-CNN是一个自然且直观的想法。但是掩码输出不同于类别和边框输出,需要提取更精细的对象空间布局。接下来,我们介绍了Mask R-CNN的关键元素,包括像素到像素对齐,这是Fast/Faster R-CNN所缺失的部分。

用于实例分割的Mask R-CNN框架

RoIAlign:虚线网格表示特征映射图,实线边框表示RoI(Region of Interest),点表示每个边框中的4个采样点。RoIAlign通过双线性插值从特征映射图上的相邻网格点计算每个采样点的值。

  • Network Architecture: 为了表述清晰,有两种分类方法
  1. 使用了不同的backbone:resnet-50,resnet-101,resnext-50,resnext-101;
  2. 使用了不同的head Architecture:Faster RCNN使用resnet50时,从Block 4导出特征供RPN使用,这种叫做ResNet-50-C4
  3. 作者使用除了使用上述这些结构外,还使用了一种更加高效的backbone:FPN(特征金字塔网络)
Head架构:我们扩展了两个现有的Faster R-CNN Head。
  • Mask R-CNN基本结构:与Faster RCNN采用了相同的two-state结构:首先是通过一阶段网络找出RPN,然后对RPN找到的每个RoI进行分类、定位、并找到binary mask。这与当时其他先找到mask然后在进行分类的网络是不同的。
  • Mask R-CNN的损失函数L = L{_{cls}} + L{_{box}} + L{_{mask}} (当然了,你可以在这里调权以实现更好的效果)
  • Mask的表现形式(Mask Representation):因为没有采用全连接层并且使用了RoIAlign,我们最终是在一个小feature map上做分割。
  • RoIAlign:RoIPool的目的是为了从RPN网络确定的ROI中导出较小的特征图(a small feature map,eg 7×7),ROI的大小各不相同,但是RoIPool后都变成了7×7大小。RPN网络会提出若干RoI的坐标以[x,y,w,h]表示,然后输入RoI Pooling,输出7×7大小的特征图供分类和定位使用。问题就出在RoI Pooling的输出大小是7×7上,如果RON网络输出的RoI大小是8*8的,那么无法保证输入像素和输出像素是一一对应,首先他们包含的信息量不同(有的是1对1,有的是1对2),其次他们的坐标无法和输入对应起来。这对分类没什么影响,但是对分割却影响很大。RoIAlign的输出坐标使用插值算法得到,不再是简单的量化;每个grid中的值也不再使用max,同样使用差值算法。

Implementation Details

使用Fast/Faster相同的超参数,同样适用于Mask RCNN

  • Training:

1、与之前相同,当IoU与Ground Truth的IoU大于0.5时才会被认为有效的RoI,L{_{mask}}只把有效RoI计算进去。

2、采用image-centric training,图像短边resize到800,每个GPU的mini-batch设置为2,每个图像生成N个RoI,在使用ResNet-50-C4 作为backbone时,N=64,在使用FPN作为backbone时,N=512。作者服务器中使用了8块GPU,所以总的minibatch是16, 迭代了160k次,初始lr=0.02,在迭代到120k次时,将lr设定到 lr=0.002,另外学习率的weight_decay=0.0001 momentum = 0.9。如果是resnext,初始lr=0.01,每个GPU的mini-batch是1。

3、RPN的anchors有5种scale,3种ratios。为了方便剥离、如果没有特别指出,则RPN网络是单独训练的且不与Mask R-CNN共享权重。但是在本论文中,RPN和Mask R-CNN使用一个backbone,所以他们的权重是共享的。(Ablation Experiments 为了方便研究整个网络中哪个部分其的作用到底有多大,需要把各部分剥离开)

  • Inference:在测试时,使用ResNet-50-C4作为 backbone情况下proposal number=300,使用FPN作为 backbone时proposal number=1000。然后在这些proposal上运行bbox预测,接着进行非极大值抑制。mask分支只应用在得分最高的100个proposal上。顺序和train是不同的,但这样做可以提高速度和精度。mask 分支对于每个roi可以预测k个类别,但是我们只要背景和前景两种,所以只用k-th mask,k是根据分类分支得到的类型。然后把k-th mask resize成roi大小,同时使用阈值分割(threshold=0.5)二值化

Experiments

Main Results

在下图中可以明显看出,FCIS的分割结果中都会出现一条竖着的线(systematic artifacts),这线主要出现在物体重的部分,作者认为这是FCIS架构的问题,无法解决的。但是在Mask RCNN中没有出现。

Ablation Experiments(剥离实验)

  • Architecture:
    从table 2a中看出,Mask RCNN随着增加网络的深度、采用更先进的网络,都可以提高效果。注意:并不是所有的网络都是这样。
  • Multinomial vs. Independent Masks:(mask分支是否进行类别预测)从table 2b中可以看出,使用sigmoid(二分类)和使用softmax(多类别分类)的AP相差很大,证明了分离类别和mask的预测是很有必要的
  • Class-Specific vs. Class-Agnostic Masks:目前使用的mask rcnn都使用class-specific masks,即每个类别都会预测出一个mxm的mask,然后根据类别选取对应的类别的mask。但是使用Class-Agnostic Masks,即分割网络只输出一个mxm的mask,可以取得相似的成绩29.7vs30.3
  • RoIAlign:tabel 2c证明了RoIAlign的性能
  • Mask Branch:tabel 2e,FCN比MLP性能更好

Bounding Box Detection Results    

  • Mask RCNN精度高于Faster RCNN
  • Faster RCNN使用RoI Align的精度更高
  • Mask RCNN的分割任务得分与定位任务得分相近,说明Mask RCNN已经缩小了这部分差距。

Timing 

  • Inference:195ms一张图片,显卡Nvidia Tesla M40。其实还有速度提升的空间,比如减少proposal的数量等。
  • Training:ResNet-50-FPN on COCO trainval35k takes 32 hours  in our synchronized 8-GPU implementation (0.72s per 16-image mini-batch),and 44 hours with ResNet-101-FPN。

Mask R-CNN for Human Pose Estimation

让Mask R-CNN预测k个masks,每个mask对应一个关键点的类型,比如左肩、右肘,可以理解为one-hot形式。

  • 使用cross entropy loss,可以鼓励网络只检测一个关键点;
  • ResNet-FPN结构
  • 训练了90k次,最开始lr=0.02,在迭代60k次时,lr=0.002,80k次时变为0.0002

MICCAI 2022:基于 MLP 的快速医学图像分割网络 UNeXt

论文地址: https://arxiv.org/abs/2203.04967

github:https://github.com/jeya-maria-jose/UNeXt-pytorch

UnetX 网络结构

Datasets

  1. ISIC 2018 – Link
  2. BUSI – Link

MICCAI 2022:基于 MLP 的快速医学图像分割网络 UNeXt

前言

最近 MICCAI 2022 的论文集开放下载了,地址:https://link.springer.com/book/10.1007/978-3-031-16443-9 ,每个部分的内容如下所示:

Part I: Brain development and atlases; DWI and tractography; functional brain networks; neuroimaging; heart and lung imaging; dermatology;

Part II: Computational (integrative) pathology; computational anatomy and physiology; ophthalmology; fetal imaging;

Part III: Breast imaging; colonoscopy; computer aided diagnosis;

Part IV: Microscopic image analysis; positron emission tomography; ultrasound imaging; video data analysis; image segmentation I;

Part V: Image segmentation II; integration of imaging with non-imaging biomarkers;

Part VI: Image registration; image reconstruction;

Part VII: Image-Guided interventions and surgery; outcome and disease prediction; surgical data science; surgical planning and simulation; machine learning – domain adaptation and generalization;

Part VIII: Machine learning – weakly-supervised learning; machine learning – model interpretation; machine learning – uncertainty; machine learning theory and methodologies.

其中关于分割有两个部分,Image segmentation I 在 Part IV, 而 Image segmentation II 在 Part V。

随着医学图像的解决方案变得越来越适用,我们更需要关注使深度网络轻量级、快速且高效的方法。具有高推理速度的轻量级网络可以被部署在手机等设备上,例如 POCUS(point-of-care ultrasound)被用于检测和诊断皮肤状况。这就是 UNeXt 的动机。

方法概述

之前我们解读过基于 Transformer 的 U-Net 变体,近年来一直是领先的医学图像分割方法,但是参数量往往不乐观,计算复杂,推理缓慢。这篇文章提出了基于卷积多层感知器(MLP)改进 U 型架构的方法,可以用于图像分割。设计了一个 tokenized MLP 块有效地标记和投影卷积特征,使用 MLPs 来建模表示。这个结构被应用到 U 型架构的下两层中(这里我们假设纵向一共五层)。文章中提到,为了进一步提高性能,建议在输入到 MLP 的过程中改变输入的通道,以便专注于学习局部依赖关系特征。还有额外的设计就是跳跃连接了,并不是我们主要关注的地方。最终,UNeXt 将参数数量减少了 72 倍,计算复杂度降低了 68 倍,推理速度提高了 10 倍,同时还获得了更好的分割性能,如下图所示。

UNeXt 架构

UNeXt 的设计如下图所示。纵向来看,一共有两个阶段,普通的卷积和 Tokenized MLP 阶段。其中,编码器和解码器分别设计两个 Tokenized MLP 块。每个编码器将分辨率降低两倍,解码器工作相反,还有跳跃连接结构。每个块的通道数(C1-C5)被设计成超参数为了找到不掉点情况下最小参数量的网络,对于使用 UNeXt 架构的实验,遵循 C1 = 32、C2 = 64、C3 = 128、C4 = 160 和 C5 = 256。

TokMLP 设计思路

关于 Convolutional Stage 我们不做过多介绍了,在这一部分重点专注 Tokenized MLP Stage。从上一部分的图中,可以看到 Shifted MLP 这一操作,其实思路类似于 Swin transformer,引入基于窗口的注意力机制,向全局模型中添加更多的局域性。下图的意思是,Tokenized MLP 块有 2 个 MLP,在一个 MLP 中跨越宽度移动特征,在另一个 MLP 中跨越高度移动特征,也就是说,特征在高度和宽度上依次移位。论文中是这么说的:“我们将特征分成 h 个不同的分区,并根据指定的轴线将它们移到 j=5 的位置”。其实就是创建了随机窗口,这个图可以理解为灰色是特征块的位置,白色是移动之后的 padding。

补充:MLP拥有大量参数,计算成本高且容易过度拟合,而且因为层之间的线性变换总是将前一层的输出作为一个整体,所以MLP在捕获输入特征图中的局部特征结构的能力较弱。通过轴向移动特征信息, Shifted MLP可以得到不同方向的信息流,这有助于捕获局部相关性。该操作使得我们采用纯MLP架构即可取得与CNN相同的感受野。

解释过 Shifted MLP 后,我们再看另一部分:tokenized MLP block。首先,需要把特征转换为 tokens(可以理解为 Patch Embedding 的过程,感觉这个就是个普通卷积,而且作者为了保证conv后的矩阵减半,设置步幅为2,总之,有些编故事的意思了)。为了实现 tokenized 化,使用 kernel size 为 3 的卷积(patch_size=3, stride=2),这样会使得矩阵H和W减半,并将通道的数量改为 E,E 是 embadding 嵌入维度( token 的数量),也是一个超参数。然后把这些 token 送到上面提到的第一个跨越宽度的 MLP 中。

这里会产生了一个疑问,关于 kernel size 为 3 的卷积,使用的是什么样的卷积层?答:这里还是普通的卷积,文章中提到了 DWConv(DepthWise Conv),是后面的特征通过 DW-Conv 传递。使用 DWConv 有两个原因:(1)它有助于对 MLP 特征的位置信息进行编码。MLP 块中的卷积层足以编码位置信息,它实际上比标准的位置编码表现得更好。像 ViT 中的位置编码技术,当测试和训练的分辨率不一样时,需要进行插值,往往会导致性能下降。(2)DWConv 使用的参数数量较少。

这时我们得到了 DW-Conv 传递过来的特征,然后使用 GELU 完成激活。接下来,通过另一个 MLP(跨越height)传递特征,该 MLP 把进一步改变了特征尺寸。在这里还使用一个残差连接,将原始 token 添加为残差。然后我们利用 Layer Norm(LN),将输出特征传递到下一个块。LN 比 BN 更可取,因为它是沿着 token 进行规范化,而不是在 Tokenized MLP 块的整个批处理中进行规范化。上面这些就是一个 tokenized MLP block 的设计思路。

此外,文章中给出了 tokenized MLP block 涉及的计算公式:

其中 T 表示 tokens,H 表示高度,W 表示宽度。值得注意的是,所有这些计算都是在 embedding 维度 H 上进行的,它明显小于特征图的维度 HN×HN,其中 N 取决于 block 大小。在下面的实验部分,文章将 H 设置为 768。

实验部分

实验在 ISIC 和 BUSI 数据集上进行,可以看到,在 GLOPs、性能和推理时间都上表现不错。

下面是可视化和消融实验的部分。可视化图可以发现,UNeXt 处理的更加圆滑和接近真实标签。

消融实验可以发现,从原始的 UNet 开始,然后只是减少过滤器的数量,发现性能下降,但参数并没有减少太多。接下来,仅使用 3 层深度架构,既 UNeXt 的 Conv 阶段。显着减少了参数的数量和复杂性,但性能降低了 4%。加入 tokenized MLP block 后,它显着提高了性能,同时将复杂度和参数量是一个最小值。接下来,我们将 DWConv 添加到 positional embedding,性能又提高了。接下来,在 MLP 中添加  Shifted 操作,表明在标记化之前移位特征可以提高性能,但是不会增加任何参数或复杂性。注意:Shifted MLP 不会增加 GLOPs。

一些理解和总结

在这项工作中,提出了一种新的深度网络架构 UNeXt,用于医疗图像分割,专注于参数量的减小。UNeXt 是一种基于卷积和 MLP 的架构,其中有一个初始的 Conv 阶段,然后是深层空间中的 MLP。具体来说,提出了一个带有移位 MLP 的标记化 MLP 块。在多个数据集上验证了 UNeXt,实现了更快的推理、更低的复杂性和更少的参数数量,同时还实现了最先进的性能。

另外,个人觉得 带有移位 MLP 的标记化 MLP 块这里其实有点讲故事的意思了。

我在读这篇论文的时候,直接注意到了它用的数据集。我认为 UNeXt 可能只适用于这种简单的医学图像分割任务,类似的有 Optic Disc and Cup Seg,对于更复杂的,比如血管,软骨,Liver Tumor,kidney Seg 这些,可能效果达不到这么好,因为运算量被极大的减少了,每个 convolutional 阶段只有一个卷积层。MLP 魔改 U-Net 也算是一个尝试,在 Tokenized MLP block 中加入 DWConv 也是很合理的设计。

代码实现:

class shiftmlp(nn.Module):
    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0., shift_size=5):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.dim = in_features
        self.fc1 = nn.Linear(in_features, hidden_features)
        self.dwconv = DWConv(hidden_features)
        self.act = act_layer()
        self.fc2 = nn.Linear(hidden_features, out_features)
        self.drop = nn.Dropout(drop)

        self.shift_size = shift_size
        self.pad = shift_size // 2

        
        self.apply(self._init_weights)

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            trunc_normal_(m.weight, std=.02)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, nn.LayerNorm):
            nn.init.constant_(m.bias, 0)
            nn.init.constant_(m.weight, 1.0)
        elif isinstance(m, nn.Conv2d):
            fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
            fan_out //= m.groups
            m.weight.data.normal_(0, math.sqrt(2.0 / fan_out))
            if m.bias is not None:
                m.bias.data.zero_()
    


    def forward(self, x, H, W):
        # pdb.set_trace()
        B, N, C = x.shape

        xn = x.transpose(1, 2).view(B, C, H, W).contiguous()
        #pad,方便后面的torch.chunk
        xn = F.pad(xn, (self.pad, self.pad, self.pad, self.pad) , "constant", 0)
        #按照dim=1维度,分成 self.shift_size(5)个块
        xs = torch.chunk(xn, self.shift_size, 1)
        #torch.roll(x,y,d)将x,沿着d维度,向上/下roll y个值
        x_shift = [torch.roll(x_c, shift, 2) for x_c, shift in zip(xs, range(-self.pad, self.pad+1))]
        x_cat = torch.cat(x_shift, 1)
        #x.narrow(*dimension*, *start*, *length*) → Tensor 表示取变量x的第dimension维,从索引start开始到(start+length-1)范围的值。
        x_cat = torch.narrow(x_cat, 2, self.pad, H)
        x_s = torch.narrow(x_cat, 3, self.pad, W)

        x_s = x_s.reshape(B,C,H*W).contiguous()
        x_shift_r = x_s.transpose(1,2)

        x = self.fc1(x_shift_r)

        x = self.dwconv(x, H, W)
        x = self.act(x) 
        x = self.drop(x)

        xn = x.transpose(1, 2).view(B, C, H, W).contiguous()
        xn = F.pad(xn, (self.pad, self.pad, self.pad, self.pad) , "constant", 0)
        xs = torch.chunk(xn, self.shift_size, 1)
        x_shift = [torch.roll(x_c, shift, 3) for x_c, shift in zip(xs, range(-self.pad, self.pad+1))]
        x_cat = torch.cat(x_shift, 1)
        x_cat = torch.narrow(x_cat, 2, self.pad, H)
        x_s = torch.narrow(x_cat, 3, self.pad, W)
        x_s = x_s.reshape(B,C,H*W).contiguous()
        x_shift_c = x_s.transpose(1,2)

        x = self.fc2(x_shift_c)
        x = self.drop(x)
        return x

class shiftedBlock(nn.Module):
    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
                 drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, sr_ratio=1):
        super().__init__()


        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
        self.norm2 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = shiftmlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
        self.apply(self._init_weights)

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            trunc_normal_(m.weight, std=.02)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, nn.LayerNorm):
            nn.init.constant_(m.bias, 0)
            nn.init.constant_(m.weight, 1.0)
        elif isinstance(m, nn.Conv2d):
            fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
            fan_out //= m.groups
            m.weight.data.normal_(0, math.sqrt(2.0 / fan_out))
            if m.bias is not None:
                m.bias.data.zero_()

    def forward(self, x, H, W):

        x = x + self.drop_path(self.mlp(self.norm2(x), H, W))
        return x