详解vmamba中的矩阵分块重排方案

点击蓝字，关注我们 本周在深入研究vmamba的代码时，发现PatchMerging2D类与PatchExpand2D类实现的很巧妙，其中还涉及到张量重排的技巧，这里做一记录。 首先总结如下，给定输入A(b,c,h,w)，h与w均为偶数。PatchMerging2D 将A的空间维度以 2X2 的不重叠邻域特征堆叠到通道维，从而得到B(b,4*c,h//2,w//2)。PatchExpand2D将B通过张量重排转换成C(b,c,h,w)。实际上在PatchMerging2D与PatchExpand2D中间还有很多步骤，这里为突出重点做出适当简化。

PatchMerging2D
如下代码所示，通过这种操作实现邻域特征提取。这里需要注意，代码只提取2*2邻域，因此PatchExpand2D中也是对应的这一情况。如果提取3*3邻域，则PatchExpand2D也将相应改变。 
        x0 = x[:, 0::2, 0::2, :]  # B H/2 W/2 C        x1 = x[:, 1::2, 0::2, :]  # B H/2 W/2 C        x2 = x[:, 0::2, 1::2, :]  # B H/2 W/2 C        x3 = x[:, 1::2, 1::2, :]  # B H/2 W/2 C                x = torch.cat([x0, x1, x2, x3], -1)  # B H/2 W/2 4*C
如果要提取3*3邻域，按照上面代码的实现逻辑，则应该是 
        x0 = x[:, 0::3, 0::3, :]  # B H/3 W/3 C        x1 = x[:, 1::3, 0::3, :]  # B H/3 W/3 C        x2 = x[:, 0::3, 1::3, :]  # B H/3 W/3 C        x3 = x[:, 1::3, 1::3, :]  # B H/3 W/3 C        ........        x7 = x[:, 2::3, 3::3, :]        x8 = x[:, 3::3, 2::3, :]                x = torch.cat([x0, x1, x2, x3,...,x8], -1)  # B H/2 W/2 9*C
PatchExpand2D
PatchExpand2D类中调用了einops.rearrange，其实现的功能为划分子块重排。 首先给出rearrange的示例，给大家一个直观感受 

x = rearrange(x, 'b h w (p1 p2 c)-> b (h p1) (w p2) c', p1=2, p2=2, c=1)
rearrange函数的输入包括三部分，x是输入张量，'b h w (p1 p2 c)-> b (h p1) (w p2) c'是张量重排规则，p1=self.dim_scale, p2=self.dim_scale, c=C//self.dim_scale是重排规则中的相关变量赋值。这个格式其实与之前我们讲过的torch.einsum很相似，详情见torch.einsum解析，但区别在于，torch.einsum实现的是爱因斯坦求和，是包含加法和乘法的。而rearrange只是张量的重排，并不涉及到加法和乘法。有人会说，那torch.einsum是不是都得有两个或以上的输入张量啊？还真不一定，比如求对角线平方和就也只需要一个输入，因此这不能作为torch.einsum和rearrange的区别。 
diagonal_elements = torch.einsum('ii->i', A)
下面我们重新看rearrange示例的重排规则部分： 
'b h w (p1 p2 c)-> b (h p1) (w p2) c'
输入为4维张量，分别为b, h, w, 2*2*1。然后将其重排为形状b, h*2, w*2, 1。我们先一句话总结其作用：将(h, w)形状的张量重排成(h*2,w*2)，也就是每个位置扩展成2*2的子区域，该区域用张量的通道维度填充，也就是这里的第四个维度2*2*1。 起初我对这一操作十分不解，核心问题是，我可以通过通过线性代数中的子矩阵概念来理解这一操作。但是我无法将张量重排操作前后的映射关系用数学公式表示出来。 例如(1,2,2,4)的张量，我们可以理解为(2,2)个向量，每个向量包含4个元素。将其通道维度展开以将原来的空间维度(2,2)扩展成(4,4)可以这样理解，原来(2,2)的每一个位置从(1,1)变成(2,2)形状，区域还是那个区域。或者反过来，(4,4)分成4份，每一份就是(2,2)。 用一张纯手图对上述文字作说明(原谅我整了个手绘就放上来了)
接下来进入解密时刻，我们还是以(1,2,2,4)->(1,4,4,1)为例。因为b=1，我们省略。拆分后的通道维数为1，也省略。初始位置(b, h, w, (p1 * p2 * c))，新位置坐标为(b, h * p1 + p1_index, w * p2 + p2_index, c) 。其中，(p1_index, p2_index)为拆分后通道维元素的相对坐标。 初始位置(0,0)包含4个元素，以第2个元素为例，其原始坐标为(0,0,1)，在拆分后的坐标为 (0,1)，对应于新位置中的(p1_index, p2_index)。这里实际上就还是按照逐行填充的原则，因为新的形状为(2,2)，所以其坐标只有{(0,0),(0,1),(1,0),(1,1)}这四种情况。这样就把映射关系的公式讲清楚了。

元素

初始位置 (b, h, w, (p1 * p2 * c))

拆分后的 (p1, p2, c)

新位置(b, h * p1 + p1_index, w * p2 + p2_index, c) 

1

(0, 0, 0, 0)

(0, 0, 0)

(0, 0 * 2 + 0, 0 * 2 + 0, 0)

2

(0, 0, 0, 1)

(0, 1, 0)

(0, 0 * 2 + 0, 0 * 2 + 1, 0)

3

(0, 0, 0, 2)

(1, 0, 0)

(0, 0 * 2 + 1, 0 * 2 + 0, 0)

4

(0, 0, 0, 3)

(1, 1, 0)

(0, 0 * 2 + 1, 0 * 2 + 1, 0)

5

(0, 0, 1, 0)

(0, 0, 0)

(0, 0 * 2 + 0, 1 * 2 + 0, 0)

6

(0, 0, 1, 1)

(0, 1, 0)

(0, 0 * 2 + 0, 1 * 2 + 1, 0)

7

(0, 0, 1, 2)

(1, 0, 0)

(0, 0 * 2 + 1, 1 * 2 + 0, 0)

8

(0, 0, 1, 3)

(1, 1, 0)

(0, 0 * 2 + 1, 1 * 2 + 1, 0)

9

(0, 1, 0, 0)

(0, 0, 0)

(0, 1 * 2 + 0, 0 * 2 + 0, 0)

10

(0, 1, 0, 1)

(0, 1, 0)

(0, 1 * 2 + 0, 0 * 2 + 1, 0)

11

(0, 1, 0, 2)

(1, 0, 0)

(0, 1 * 2 + 1, 0 * 2 + 0, 0)

12

(0, 1, 0, 3)

(1, 1, 0)

(0, 1 * 2 + 1, 0 * 2 + 1, 0)

13

(0, 1, 1, 0)

(0, 0, 0)

(0, 1 * 2 + 0, 1 * 2 + 0, 0)

14

(0, 1, 1, 1)

(0, 1, 0)

(0, 1 * 2 + 0, 1 * 2 + 1, 0)

15

(0, 1, 1, 2)

(1, 0, 0)

(0, 1 * 2 + 1, 1 * 2 + 0, 0)

16

(0, 1, 1, 3)

(1, 1, 0)

(0, 1 * 2 + 1, 1 * 2 + 1, 0)

其实这里有一个有趣的地方，如果 'b h w (p1 p2 c)-> b (h p1) (w p2) c' 被改成了 'b h w (p1 p2 c)-> b (h p2) (w p1) c' 会发生甚么事情？我们举个最简单的例子作为结束。

原来的结果如果是 

则修改后的结果是

也就是说，本来是逐行填充，现在变成了逐列填充！
尾声

实际上，我在和师兄探讨这个映射关系具体是什么的时候，师兄的意思是不求甚解，原因是这只是一个工具，我们知道怎么用即可，不必深究。其实从科研效率的角度上来讲，师兄说得完全正确。人的时间和精力是有限的，需要投入到最紧要的事情上。但我认为在rearrange这个上花费的时间是有意义的。理由如下： 其一，二维图像的分块与重排的确非常重要，例如卷积操作的滑动窗口其实本质上就是一种分块，而由于图像内目标的大小、形状等存在差异，实际上不同的分块策略会对模型性能有影响，了解张量重排的映射关系对于深入理解卷积神经网络是有益处的。 其二，目前深度学习已经发展了很多年，以发论文为例，审稿人的口味越来越刁钻，简单的增删改模块已经难以发表好论文。我认为想要有所收获，一方面是深度，需要对原理有更深入的理解，从而有所改良；另一方面是广度，需要旁征博引，他山之石可以攻玉。

综上所述，科研本就是需要不断学习，你也不知道今天学的东西能否用上，这个确实是玄学。但是如果有兴趣有精力有时间，还是要深入探讨一下，否则以后不管是进入业界做项目还是进入高校教书育人，对本领域的知识还是得学学许昕，“还是太全面了” 终归是好词。

END 你 你 好 春 天 迟日江山丽 春风花草香

©著作权归作者所有,转载或内容合作请联系作者
  • 人面猴
    序言：七十年代末，一起剥皮案震惊了整个滨河市，随后出现的几起案子，更是在滨河造成了极大的恐慌，老刑警刘岩，带你破解...
    沈念sama阅读 205,132评论 6赞 478
  • 死咒
    序言：滨河连续发生了三起死亡事件，死亡现场离奇诡异，居然都是意外死亡，警方通过查阅死者的电脑和手机，发现死者居然都...
    沈念sama阅读 87,802评论 2赞 381
  • 救了他两次的神仙让他今天三更去死
    文/潘晓璐 我一进店门，熙熙楼的掌柜王于贵愁眉苦脸地迎上来，“玉大人，你说我怎么就摊上这事。” “怎么了？”我有些...
    开封第一讲书人阅读 151,566评论 0赞 338
  • 道士缉凶录：失踪的卖姜人
    文/不坏的土叔 我叫张陵，是天一观的道长。 经常有香客问我，道长，这世上最难降的妖魔是什么？ 我笑而不...
    开封第一讲书人阅读 54,858评论 1赞 277
  • 港岛之恋（遗憾婚礼）
    正文 为了忘掉前任，我火速办了婚礼，结果婚礼上，老公的妹妹穿的比我还像新娘。我一直安慰自己，他们只是感情好，可当我...
    茶点故事阅读 63,867评论 5赞 368
  • 恶毒庶女顶嫁案：这布局不是一般人想出来的
    文/花漫 我一把揭开白布。 她就那样静静地躺着，像睡着了一般。 火红的嫁衣衬着肌肤如雪。 梳的纹丝不乱的头发上，一...
    开封第一讲书人阅读 48,695评论 1赞 282
  • 城市分裂传说
    那天，我揣着相机与录音，去河边找鬼。 笑死，一个胖子当着我的面吹牛，可吹牛的内容都是我干的。 我是一名探鬼主播，决...
    沈念sama阅读 38,064评论 3赞 399
  • 双鸳鸯连环套：你想象不到人心有多黑
    文/苍兰香墨 我猛地睁开眼，长吁一口气：“原来是场噩梦啊……” “哼！你这毒妇竟也来了？” 一声冷哼从身侧响起，我...
    开封第一讲书人阅读 36,705评论 0赞 258
  • 万荣杀人案实录
    序言：老挝万荣一对情侣失踪，失踪者是张志新（化名）和其女友刘颖，没想到半个月后，有当地人在树林里发现了一具尸体，经...
    沈念sama阅读 42,915评论 1赞 300
  • 护林员之死
    正文 独居荒郊野岭守林人离奇死亡，尸身上长有42处带血的脓包…… 初始之章·张勋 以下内容为张勋视角 年9月15日...
    茶点故事阅读 35,677评论 2赞 323
  • 白月光启示录
    正文 我和宋清朗相恋三年，在试婚纱的时候发现自己被绿了。 大学时的朋友给我发了我未婚夫和他白月光在一起吃饭的照片。...
    茶点故事阅读 37,796评论 1赞 333
  • 活死人
    序言：一个原本活蹦乱跳的男人离奇死亡，死状恐怖，灵堂内的尸体忽然破棺而出，到底是诈尸还是另有隐情，我是刑警宁泽，带...
    沈念sama阅读 33,432评论 4赞 322
  • 日本核电站爆炸内幕
    正文 年R本政府宣布，位于F岛的核电站，受9级特大地震影响，放射性物质发生泄漏。R本人自食恶果不足惜，却给世界环境...
    茶点故事阅读 39,041评论 3赞 307
  • 男人毒药：我在死后第九天来索命
    文/蒙蒙 一、第九天 我趴在偏房一处隐蔽的房顶上张望。 院中可真热闹，春花似锦、人声如沸。这庄子的主人今日做“春日...
    开封第一讲书人阅读 29,992评论 0赞 19
  • 一桩弑父案，背后竟有这般阴谋
    文/苍兰香墨 我抬头看了看天上的太阳。三九已至，却和暖如春，着一层夹袄步出监牢的瞬间，已是汗流浃背。 一阵脚步声响...
    开封第一讲书人阅读 31,223评论 1赞 260
  • 情欲美人皮
    我被黑心中介骗来泰国打工， 没想到刚下飞机就差点儿被人妖公主榨干…… 1. 我叫王不留，地道东北人。 一个月前我还...
    沈念sama阅读 45,185评论 2赞 352
  • 代替公主和亲
    正文 我出身青楼，却偏偏与公主长得像，于是被迫代替她去往敌国和亲。 传闻我的和亲对象是个残疾皇子，可洞房花烛夜当晚...
    茶点故事阅读 42,535评论 2赞 343

推荐阅读更多精彩内容