本文直接采用VanillaNet的方案重新开始构建模型ReduceNet。按照VanillaNet的方式极限压缩网络深度至单层

前几天在知乎上看到华为提出的一个VanillaNet,其中的一个设计点和我一直想实现的功能非常类似,即训练阶段的时候模型是比较深的网络,推理的时候会自动变成比较浅的网络。起初我一直想缓渐地暴力地丢掉卷积层,结果不能如愿,模型性能会在训练阶段的最后几个epoch一直退化,模型变成毫无用处的废物。整整大半年一直也没找到稳定性能的方案,这也把我后续所有的构想都否定了。本着拿来主义,本文直接采用VanillaNet的方案重新开始构建模型ReduceNet(reduce depth,还有就是张量的那个reduce)

ReduceNet和VanillaNet有以下几个区别:

1. a) 引入残差,提升模型性能。部署阶段,残差也可以作为特殊的卷积核融入卷积算子

b) 除了用一个VanillaNet中的"可消解的激活函数"外,不使用多余的非线性激活函数,从而使得不同模块继续融合。这也是reducenet关键设计,起源于某个偶然的实验发现:resnet两个3x3卷积构成的模块,去掉第二个激活函数,性能会提升,顺便提高了速度。所以reducenet模块也只采用一个激活函数,而且是在训练结束后会消失的激活函数

c) 不使用maxpooling,用卷积代替,不同stage可以继续融合

2.模型只有在训练阶段存在激活函数,部署阶段整个网络没有任何激活函数(最后的softmax还是需要的)

3.根据2,backbone所有的线性层理论上可以继续融合为一个超大尺寸卷积核的卷积层,逼近一个全连接层,印证了那个传说中的“单层网络理论”(印象中是有的,貌似是单层大网络学习容量足够却很难拟合之类的,我不记得出处了,有知道的可以告知一下)。

4. 利用bottleneck结构增加模块网络宽度,为训练阶段引入更多参数量,而不增加推理时候的开销。理论上我们可以在训练时让模块中间无限宽,推理时计算开销只由模块首尾的宽度决定。实际上应该会存在短板效应。我们大概只能在首尾宽度固定情况下不断压榨性能直到性能饱和。再继续增加中间网络宽度估计就变成冗余了。这时只能增加模块首尾宽度继续提升性能。用这种方式有点神经架构搜索的味道。也许神经神经架构搜索本身是很冗余的领域。渐进地调整宽高,压榨性能,试探临界才是较优解。

本文贡献如下:

1.理论说明深度非线性网络可以融合成单层网络,实验中模型也取得不错的性能。本文提供了基本模块的思路,目前的SOTA模型只要对应魔改,在不改架构的情况下,尽情利用本文模块的特性榨取模型性能就可以,理论上100%涨点,不增加额外推理开销。其实,本文的东西是提供了“填充”的功能,任何架构都可以通过本文模块中“无限”和“缩并”的特性,理论上压榨出该架构下最高性能。所以在数据充足的情况下,结合本文的模块,最终的模型瓶颈还将会是架构设计。不过目前为止,我们还远远还没摸到瓶颈。暂时只要用sota架构继续压榨就可以了。本人资源有限,没卡,只要负责验证方法的有效性就可以了。

上面说的是如何利用现有架构提升性能。另外我这里提供个模型设计方案去探索新的适合应用场景的sota架构设计方案。我们先定下模型的开销约束,比如计算量,参数量,内存占用。在符合开销约束情况下,扫描网络深度和每层的宽度(这里是指部署的宽深,训练阶段,一层网络是可以根据我们的办法无限扩张宽深的,只要最后缩并回来),这样就有不同的架构,让它们一一被压榨性能就可以了,因为宽深固定,性能极限就固定了。和SOTA比较,胜者为王。和神经架构搜索相比,我们可以做到专心压榨性能,相当于搜索得更快了。因为神经架构搜索方法是没法保证一定宽度深度各种组合下性能如何的,再加上搜索空间这么大,很容易三心二意顾此失彼。

总之,有一个证明,两个pipeline方案

2.尝试解释了为什么单层网络规模足够大却很难获得学习能力,从优化的角度为多层非线性网络和单层网络建立了联系,并指出多层非线性网络本质上可以为"探索单层网络的权重最优值”提供更大的探索自由度, 缓解优化算法和单层网络的“不和”,从而具备更高的学习效率。

3.训练阶段,ReduceNet可以和现有的其他模型设计兼容,享有深度非线性网络的性能收益。推理阶段则最终可以变成单层网络,保持高性能的同时,获得激进的模型效率。当然,还有一个好处是,部署网络的形态是具备灵活性的,你爱融合几层就几层,依照硬件的内存占用,延时条件调整融合的层数。比如,两层3x3可以变成5x5卷积,在某个硬件上可能就能提速了。如果任务很容易,很小的模型就能满足,那直接压成单层,在硬件合适的情况下,推理速度会非常快,提供了用深度网络高效学习,用单层网络快速推理的完美方案。

4.ReduceNet理论上可以在训练阶段无限深和无限宽(bottleneck中间表示无限宽),推理阶段变成一个小网络,有需要的话甚至直接变成单层网络。

不排除我犯了很傻的代码错误,一切只是我的妄想(真的碰到好几次了),如果错了就删帖跑路

代码链接

[https://github.com/ohmydroid/reducenet

灵感来源

模型压缩的核心思想是先利用大网络的学习能力,最后通过各种手段(知识蒸馏,剪枝,神经架构搜索,动态推理网络等)将该能力转移到较小的模型上,以减少资源开销。(目前的分类习惯并不把神经架构搜索,动态推理网络归于模型压缩,这仅仅是我个人习惯和理解)

ReduceNet的灵感来源有以下几个:

1.RepVGG是比较特别的模型,训练时有多个分支,推理时可以合并成单个分支,可以看作是网络宽度这个角度的缩并。所以很容易联想到把这种设计扩展到网络深度方向,相关的工作有DepthShrinker,RMNet。严格来说,读研期间设计的静态condconv比这个更早,模块化动态卷积又被condconv先发。横向重参化的idea应该来源于condconv.

2.SkipNet, 一种动态推理网络,在推理阶段通过路由网络为不同数据选择性地跳过主网络的一些网络层,以减少计算.

3.Stochastic Depth,训练阶段以概率的方式激活或者不激活残差块中的卷积分支,可以提升性能,但是推理阶段没有减少网络深度,按照概率融合卷积块和残差。

4.DepthShrinker, 比较接近我设想中的模型,但是条件很麻烦,需要学习mask参数,还要蒸馏,微调和融合。RMNet需要剪枝,微调,和融合。这些方式其实和华为的VanillaNet比较接近了。后者更加简洁有效。我自己想的方案暂时又粗暴且无用,打不过就加入。

5.VanillaNet中Deep Training Strategy,里面的公式以及衰减策略和我设计的几乎差不多,只是VanillaNet选择了对激活函数下手,再后期融合两个串联的卷积算子。而我天真地以为idea结合了skipnet,Stochastic Depth,神经架构搜索等思想必定万无一失(不排除以后会有正确稳定的方案),从而贪心地选择最后丢弃整个卷积层。实验结果一度让我怀疑人生,模型性能会在最后阶段彻彻底地崩溃。换成VanillaNet的Deep Training Strategy后模型结果非常理想。

6. 是我N年前构想的DetachNet,算是ReduceNet的雏形,希望有一个无限大的网络可以为一个小网络的每一层提供scale factor(scale的方式乘,加,卷积都可以)。这个和动态卷积非常类似了,只不过我希望大网络在训练的可以无限大,推理的时候可以摘除,将能力转移给小网络,完美的妄想的无限压缩算法。最终败于BP算法。也不是不可能。我一度相信,这种模型会让AGI时代到来,完全没想到AGI风暴是ChatGPT引起的。ReduceNet现在好像也可以在训练阶段无限大,推理阶段非常小,所以DetachNet也算后继有net。

压缩深度方法

ReduceNet_子网

总的来说,这种设计使得模型在训练过程中借助多层卷积和非线性获得更强的学习能力;推理阶段,非线性消失则为减少网络深度提供了可能。

顺便一提,VanillaNet网络虽然浅但是非常宽,计算量,参数量,内存占用都很大,唯一的优势可能就是用更快的速度取得接近sota的准确率。然而,这速度也是有限制的。归根结底这限制来自于宽网络引发的一系列连锁反应。庞大的通道数,会带来更多的内存占用,进一步增加访存。batch size等于1的时候vanillanet还有速度优势,其远超深网络的计算开销还能忽略。如果batch size变大呢,想必速度优势不会有这么明显,甚至会逐渐慢于普通的sota模型,那其优势将荡然无存。

VanillaNet的deep training strategy足以对深度网络模型设计产生深远的影响,可能在网络架构上设计元素太多太杂,反而不能发挥整体优势。

ReduceNet_子网_02

模块设计 

ReduceNet_卷积_03

和VanillaNet不同,本文设计的基本模型由3x3 conv, 1x1 conv ,残差连接构成。下采样的话,也只要把3x3 conv的stride设置为2,并且去除收尾的残差(激活函数的残差保留)。值得注意的是整个模块只有一个激活函数,训练结束,激活函数也会消失。之所以只用一个激活函数是因为之前发现resnet的残差块中去掉一个激活函数效果会更好,速度还更快(查重发现被人发过了)。

ReduceNet_子网_04

训练阶段,模块如左图,推理阶段,这个模块会整合成一个纯线性的3x3 conv。能够融合的原因有以下几点:

1.conv,bn可以融合成1个conv

2.一个3x3 conv和一个1x1 conv会整合成一个3x3 conv

3.残差连接是可以看作一种特殊的卷积核,所以也可以参与线性融合

比起VanillaNet,ReduceNet的其中一个优势在于能够和现有的架构设计方式兼容。

此外,当训练结束后,所有算子完成上述的融合后,ReduceNet的网络backbone部分,只有连续的3x3 conv,这意味着这些纯线性的3x3 conv可以进一步连续融合直至变成一个更大尺寸卷积核的单层卷积。最终,整个网络会变成一个卷积层,pool, FC, 总共3层。pool是avgpool的话网络应该可以直接变成1层。

实现效果

按照resnet56 cifar的架构,reducenet56(训练阶段是56层) cifar10的准确率可以达到93%。resnet的模块是两个3x3 conv, 我的模块是一个3x3 conv+一个1x1 Conv(可以融合成一个3x3 Conv)。

如果只融合模块里的卷积不继续融合的话,网络深度应该是29. 为此构造了一个29层的ResNet进行比较(和原版resnet不同,每个模块是一个非线性3x3卷积加一个残差,因为要和我的模型推理阶段的架构公平对比)准确率是91+%(手贱,不小心把tmux退出了)。具体架构和训练代码可以参考repo的代码。虽然只是在cifar10上的实验,但是效果已经足够惊艳。代码里还没有加入任何融合操作,我就是确认纯线性条件是不是成立,能跑不错的结果就行,结果超出预期。

如果设计逻辑正确无误,代码也能确保线性条件成立的话,可以印证:

1.单层网络的表达能力是足够的,只是我们目前的优化算法并不能让单层网络找到合适解。不是模型容量的问题,是目前深度学习的优化算法碰上单层线性网络的天然局限。

2.深度非线性网络最终是可以转换成单层网络的,多层非线性网络本质上可以看作是为单层网络高效搜索权重的方法,弥补了当前单层网络优化算法的缺陷。

解释和猜想

这么多的非线性层也许只是为单个线性层找到更好的参数。最开始,参数只是随机的,非线性的存在和BP算法机制,让整体网络的每个参数个体能够像布朗运动一样“更自由”地改变参数。更确切地说,是随机初始化让前向传播有了一定随机性,比如ReLU后的数值,随机地为0。因果循环,这种随机的失活,又让反向传播只能随机地寻找路径让梯度流动(后面会提到,这是天然的子网采样,每个子网都有机会接受数据洗礼)。

训练阶段,堆叠纯非线性等价于一层线性层,即便堆叠再多,所有参数的更新轨迹都被“限定”了,训练前的初始值给整个模型留下了深刻的烙印并且贯穿整个训练过程。因为无论怎么更新,初始值已经决定了反向传播梯度分配的比例,W这个整体权重完全散失了“随机探索超流形的自由性”。

深度非线性网络和巨大的参数量,或许只是为优化过程买更多的彩票。参数越多,拓扑路径组合的子网络越多,W就能更加细致地,发挥群体力量更加高效地探索空间,最终只是为最后的单层网络权重服务(这个过程当然也很有可能造成冗余。这种“冗余”某种程度上是有意义的,可以用计算开销换取更快的学习速度,有粒子群优化的味道)改日再继续想这个优化问题,因为是个很大的数学问题,而我只是凭借文字描述抓住一点鳞片,暂时无法更确切地描述。单单在多层非线性网络和单层网络之间建立辨证的数学联系,这本身就是意义深远的事情。假如深度非线性网络的存在本身就是在自由随机地为单个线性层(没有考虑softmax)寻找"无数"可能性,网络最终可以缩并成一层,那这实在是太具有暴力美学的数学哲思。应该可以颠覆现在优化领域的一些常识,也可能可以开个新的方向。

这并不是说模型可以无限压缩,只有在深度非线性网络学足以学到东西的时候,我们才会继续转换成单层网络。单层网络开销有多大还是得由当前的深度非线性网络所决定。所以,如何设计高效的深度网络仍然是一个重要的课题,剪枝、蒸馏、神经架构搜索,动态推理网络,这些方向或许本质上被某个共同的"规则"所左右,未来可能统一,神经架构搜索(NAS)可能会被废弃(因为实在是太浪费了,NAS本质上是一种高度浪费的压缩算法)。如果说这些算法实在压缩,那么还有必然还一种对立的,是生长型的算法。小模型性能饱和之后,可以继续成长成更大的网路从而继续压榨更高的性能。明明是对等的产物,不知道为什么少有人问津。

NAS和剪枝简单讨论

NAS设置的超网搜索空间实在是太大了,打个比方,就像把食物浅尝几口就丢掉。获得超网也并不需要这么多并行的层。最简单的VGG本身就是一个超网。我的意思是说,当训练阶段仔细探索子网,整个模型本身相对来说就是一个超网。与其尝一口,嚼几下就丢掉,还不如仔细咀嚼每一口。换言之:

1.需要充分地让子网接受数据训练的洗礼

2.用有限的计算资源在最简单的网络中创造更多数量的子网。当充分训练,子网便会“自动集成”成高性能的模型。

当有了足够的子网,便意味着有了足够的学习容量。但这还不够,每一个子网(它们之间也会共享部分,难分彼此)需要“相对独立”地进行学习。每个子网地能力足够强,整体网络性能才会更高。我们并不需要额外的操作,像ReLU这种网络就是天然的子网采样器。训练完毕,网络本身就是所有子网的集成。

上面说过,当所有子网模型都可以独当一面的时候,这就意味着整体网络足够健壮。那些剪枝算法只是回过头来再选取合适的子网,使其成为更小的独立模型。

所以有时候,我们大概不能片面地评估一个剪枝算法一定比另一个算法更好。因为当训练不充分,某些剪枝算法也会因为运气好(不好)(也是随机)采样到一些优秀(糟糕)的子网,这样的比较并没有太大的意义。除非它们自己能够证明自己是“精准定位”到优秀子网的。举个例子作类比,能不能打到好的猎物取决于山林资源条件如何,猎人身手如何。在贫瘠的山林里,再强的猎人也不见得收获会比菜鸟猎人更丰。明显只有资源足够丰富的情况下,猎人的身手才能成为关键。当然也会存在一种极端情况,山林资源过于丰富,再怎么菜鸡的猎人闭着眼睛乱射都能收获满满。这也对应下文所说的,随机剪枝完全足够,不逊色任何算法。

原则上应该让随机剪枝大量采样子网,证明网络的健壮性。再让其剪枝算法进行比较(比较采样的良率,搜索时间,性能,效率等)。神经架构搜索(NAS)本身就可以看作是剪枝算法,但是大多数情况下NAS的超网太大了,不能也根本没必要充分训练。再回过头来看,如果我们已经让一个模型具备良好的健壮性了,花里胡哨的剪枝算法也就没有什么太大的意义了(如果是那种剪枝,微调最后还不和相同规模的模型继续比较那无异于掩耳盗铃了。而且剪枝的网络比大网络更强也并不奇怪,毕竟随机初始化,学习率主场都可能是关键。没有统计意义的胜负也没什么参考价值。)这个时候,随机剪枝就足够了,考虑到模型效率,必须是按照规则形状的随机剪枝。

VGG这种超网也可以变得强大。比如引进残差,vgg的子网数量就变得更多了。比如随机深度算法,让一些更小的子网能够独立且充分地接受数据洗礼。个人认为上述两条设计原则应该可以为轻量化模型设计提供一定程度的方向性指导。超网和子网的辨证关系想必对深度集成理论会有所帮助。

此外,如果模型能被压缩(剪枝,蒸馏等),性能还不下降,说明大模型本身性能不饱和。也可以说是大模型存在结构冗余。现实中,我们可以利用大模型快速收敛得到不错的性能,再想方设法利用各种奇技淫巧压缩模型,小心翼翼地防止精度损失。其实这和本文提出地bottleneck缩并技术很类似,都要经历膨胀再收缩的过程。但是我觉得bottleneck(层积,模块级,网络级都可以)缩并技术更加简洁,两端网络宽度固定,以满足资源约束,中间尽可能大以获得强大的学习能力,大到性能饱和为止,最后缩并成很窄的网络即可,不需要花里胡哨的过程。这个貌似直接解决了几年前所谓的信息瓶颈理论。(我不记得它说了什么,听名字和bottleneck 缩并应该还是很相关的。应该不需要分析,直接做出来即可)。

Bottleneck 缩并技术实验

做个实验就很明白了。目前repo实验代码用的是layer级别的bottleneck设计,当然完全可以选择block级,甚至整个网络是一个bottleneck. 理论上net级容量是最大的,可以更快压榨性能。

ReduceNet20, 调整expansion可以控制bottleneck中间的宽度,3个模型最终部署成本都是一样的,简直是免费的性能榨汁机。我还没有探索expansion能有多大,有时间补上数据。20层只是训练层数,之前说过了,部署的网络形态是可以灵活改变的。比如,至少1x1卷积就可以消掉,网络变成11层(10个3x3 conv, avgpool,FC),剩下3x3如果融合就要考虑计算量的增加自行trade-off了。硬件条件合适的话,计算量就能换更高速的模型推理。

ReduceNet_激活函数_05

代码实现和讲解

完整训练代码:https://github.com/ohmydroid/reducenet

如果是我搞错了,别打我。发这种半成品实在是因为已经浪费了好几个月在研究貌似没有什么用的东西。

模块代码如下:

class BasicBlock(nn.Module):

    def __init__(self, in_planes, planes, stride=1,scaler=0.,expansion=1):
        ## expansion factor allows bottleneck structure for basi block. There exists high dimensional intermediate representation during training, but only low dimensional representation will be left during inference.
        
        super(BasicBlock, self).__init__()

        self.shortcut = True if stride==1 else False
        self.scaler = scaler

        self.conv1 = nn.Conv2d(in_planes, expansion*planes, kernel_size=3, stride=stride, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(expansion*planes)

        self.conv2 = nn.Conv2d(expansion*planes, planes, kernel_size=1, stride=1, padding=0, bias=False)
        self.bn2 = nn.BatchNorm2d(planes)


    
    def forward(self, x):
        out = self.bn1(self.conv1(x))
        out = self.scaler*F.relu(out, inplace=True) + (1-self.scaler)*out
        # out = nn.Dropout2d(drop_rate)(out), to be tested
        out = self.bn2(self.conv2(out))
        if self.shortcut:
           out += x 
        return out

    '''    
    def forward(self, x):
        out = self.bn1(self.conv1(x))
        #out = self.scaler*F.relu(out, inplace=True) + (1-self.scaler)*out
        #out = self.bn2(self.conv2(out))

        if self.shortcut:
           out = out+x 
        
        out = F.relu(out,inplace=True)
        return out
    '''
  • 1.
  • 2.
  • 3.
  • 4.
  • 5.
  • 6.
  • 7.
  • 8.
  • 9.
  • 10.
  • 11.
  • 12.
  • 13.
  • 14.
  • 15.
  • 16.
  • 17.
  • 18.
  • 19.
  • 20.
  • 21.
  • 22.
  • 23.
  • 24.
  • 25.
  • 26.
  • 27.
  • 28.
  • 29.
  • 30.
  • 31.
  • 32.
  • 33.
  • 34.
  • 35.
  • 36.
  • 37.
  • 38.
  • 39.

可以看到,模块中只有一个激活函数,即

out = self.scaler*F.relu(out, inplace=True) + (1-self.scaler)*out
  • 1.

self.scaler 就是一个不可训练,人为控制数值的参数 \lambda\lambda 。在训练过程中,随着iteration(每个mini batch)按照cosine方式从1衰减至0。训练结束后,这个激活函数就会退化成identity mapping(输入等于输出)。那么conv1,bn1,conv2,bn2四层可以融合成单层的3x3 conv。之后,把模块首尾的残差变成3x3卷积核的形状,"对角线"元素填充为1(这比1x1 卷积核形状的残差稍微复杂一点点,不是正规的对角线,是3x3的中心点,每个通道沿着对角线,可以参考DiracNet),和3x3卷积继续横向融合。因为推理阶段不存在任何非线性函数,所以所有的线性算子横向和纵向都可以融合,这里不再赘述。

模块中self.scaler可以由ReduceNet的Module的参数传入。main.py文件控制self.scaler的数值,cosine衰减。代码如下:

iters_per_epoch = len(trainloader)

    for batch_idx, (inputs, targets) in enumerate(trainloader):
        
        prop = ((epoch-1)*iters_per_epoch+batch_idx+1)/(args.epoch*iters_per_epoch) 
        scaler = torch.tensor(0.5*(1+np.cos(np.pi*prop)))

        inputs, targets = inputs.to(device), targets.to(device)
        #print('scaler is {}'.format(scaler))
        net.scaler.data=scaler
        optimizer.zero_grad()
        outputs = net(inputs)
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()
  • 1.
  • 2.
  • 3.
  • 4.
  • 5.
  • 6.
  • 7.
  • 8.
  • 9.
  • 10.
  • 11.
  • 12.
  • 13.
  • 14.
  • 15.
未来的工作

以后继续优化:

1.和RepVGG等重参结合

2.3x3卷积可能不需要这么多,3x3卷积的数量关乎模型性能和效率,怎么寻找合适的数量值得研究

3.最大的愿望当然是能够想出一步到位的方案,而不是用VanillaNet这种策略

4.还没有写融合算子和整个网络的代码

5. 尝试在conv旁边多加一条残差

6.参考BagNet, MLP-mixer, weighted sum of patches is all you need

7.我比较好奇的是究竟能不能根据训练结果的实时反映生成一个控制 \lambda\lambda 的信号,再让这个新的数值再控制网络形成控制闭环。

8.拥抱LORA。初闻LORA的时候我就意识到了,这东西和我从前想的多核卷积太像了,与我设想的无限可缩并的bottleneck不谋而合。总的来说,这种类似“多重宇宙”,“影分身般”的特性无处不在。比如,谷歌某个集成方法,训练多个模型,最后用加法融合每个模型的参数,最后微调整。不过这个方法也可以看作上文所说的net级别的bottleneck缩并技术的一个稀疏子集。

9.最难的恐怕是我该如何用一套数学语言去严谨地解释多层非线性网络是在寻找单层网络的权重这一动态过程。

10.融合后可以继续在现有网络(训练过的ReduceNet)基础上生出新的活性网络,可以是横向的扩展宽度,纵向地插入新层,甚至每个模块能成为bottleneck中间的小网络,成为被缩并的对象,成就新的下一代模块。总而言之,继承旧网络,在此基础上附上新网络继续扩展学习容量,压缩融合,嫁接生长,循环反复。

代码模块其实和resnet没什么两样,可以自行验证相同架构网络的效果。把conv1换成conv3,就可以和现有的resnet架构完全兼容。

最后解释一下免得让人误解。ReduceNet在训练中理论上可以无限深无限宽,但这并不意味着就模型最后一定具有无限的学习能力。部署时模型模型(都统一为单层网络比较好了)的宽度等物理因素决定了它的能力上限,训练阶段无限的资源投入都只是让最终的模型无限逼近固有的能力上限。这个规律同样适用于现在的大模型。没有免费的涌现。计算够不够久,数据够不够多,都是瓶颈,能量换信息所有过程都是为了逼近这个上限。突然对涌现这个词语这么疯魔,只是低估了这个上限了而已。