强化学习应用到模型压缩

本文介绍几篇自动化模型压缩相关论文的主要思想,包括网络结构搜索及自动化模型压缩压缩

一、强化学习基本思想回顾

(一)、简介

强化学习是一种无监督学习方法,数据集不存在标注信息。而是通过环境和机器之间的“互动”,让机器获得一种最优化策略,在这种策略驱动下,环境的不同状态引发机器的最优动作。

用X表示环境的状态,a表示在这样的状态下机器所采取的动作,采取了这样的动作后环境会反馈给机器一个奖励r。机器学习到使得奖励最大化的策略。

通常环境被建模成一个马尔科夫过程,即机器采取的动作只决定于环境当前的状态,和过去、将来的状态无关。也即是机器产生的动作a导致环境状态转移到新的状态,这个状态只和环境前一个状态有关。

(以下引自周志华《机器学习》)暂且假定任务对应的马尔科夫决策过程四元组E=<X, A, P, R>均为已知,这样的情形称为“模型已知”,即机器对环境进行了建模,能在机器内部模拟出与环境近似的状况。

对于任意的状态x,x’和动作a,在状态x下执行动作a转移到x’的概率为Pax-x’,转移带来的奖赏为Rax-x’。定义一个累积奖赏函数:


image.png

强化学习要做的工作就是最大化这个奖赏函数:


image.png

这个就是最优策略。在这样的策略下,给定一个环境状态,机器就可以采取最优的动作。可以将以上公式描述成Bellman等式,这个等式用于强化学习最优策略求解。
image.png

其实我们仔细观察强化学习,发现它和有监督学习优点类似。如果把状态空间作为数据集,奖励作为标注,那么机器也是在进行分类。只不过现在这些标注并不是提前知道的,而是在最后的结果给出来的。强化学习在某种意义上可以看做具有“延迟标记信息”的监督学习问题。(引用字周志华《机器学习》)

image.png

(二)、常用强化学习算法

1、DDPG

(1)、简介

DDPG是为了解决DQN算法只能用于离散行为空间问题而产生的。

DQN用一个DNN网络来拟合Q函数(评价累积奖励的函数),它的输入是当前环境状态,而输出是多个离散行为。其训练过程中使用DNN网络的Q值和目标Q值的平方差作为损失函数。如果一个任务有多个行为,那么DQN可以通过采用更细分的离散动作来处理。但是会导致巨大的计算量。

(2)、内容

DDPG解决了动作连续的问题。DDPG依赖于actor-critic结构。
Actor用于调节策略函数的参数,并能输出最终动作。

image.png

Critic是一个评价结构,用于评估策略函数。它使用一个叫做临时差异错误率作为评估函数:

image.png

整个过程如下图:
image.png

(3)、TIPS

为了解决连续性动作空间的问题,DDPG在action基础上增加了noise。一种方式是在输出的动作上直接加noise,另外一种是在DNN网络中加noise。


image.png

发明者认为在参数空间加noise效果更好。

(4)、算法流程

对策略函数的搜寻是通过更新策略中参数来实现的。即通过梯度下降方法来更新。假设J是一个整体表现函数,那么其梯度为:


image.png

算法如下:


image.png

1、REINFORCE

这个算法有一个网络agent,agent网络中每个单元的输入或者来自于环境或者来自于前几个单元。其实也就是一个DNN网络结构。每个单元输出是随机化的。假设一个单元参数为w_ij为y_i,输入为x_j。那么用一个概率函数来描述输出为:


image.png

比如如果使用Bernoulli分布,这个概率函数就是:


image.png

如何来更新网络中参数呢?也是采用随机策略梯度的方式:
image.png

其中有:


image.png

二、剪枝AutoML for Model

Compression(AMC)
https://zhuanlan.zhihu.com/p/296530361
AutoML for Model Compression(AMC)是一个利用强化学习自动搜索并提高模型剪枝算法质量的框架,其完整的流程如下:

image.png

在一般的剪枝算法中,我们通常遵循一些基本策略:比如在提取低级特征的参数较少的第一层中剪掉更少的参数,对冗余性更高的FC层剪掉更多的参数。然而,由于深度神经网络中的层不是孤立的,这些基于规则的剪枝策略并不是最优的,也不能从一个模型迁移到另一个模型。

AMC方法便是在该背景下,利用强化学习自动搜索并提高模型压缩的质量,该框架是每一层进行独立压缩,前一层压缩完之后再往后层进行传播,t层接受该层的输入特征s_t,输出稀疏比率a_t,按照a_t对该层进行压缩后,智能体移动到下一层L_t+1,使用验证集精度作为评估,伪代码如下:

image.png

0、主要思想

采用DDPG作为agent,用于决定压缩网络的具体举措。也就是强化学习中所说的动作。环境是待压缩网络,包括网络的各种参数。它对网络逐层压缩,通过网络反馈的状态s_t,给出这层网络的压缩系数a_t。完成压缩后再转移到下一层网络进行。

1、状态空间

每层网络有11个参数来表示其状态:


image.png

其中

  • t是层的编号,
  • 卷积核的维度为nxcxkxk,其中n为输出通道,c为输入通道。对网络压缩可以减小其nxc维度。
  • Stride是卷积运算的步长。
  • FLOPs是每层的计算复杂度,即这层卷积需要多少次运算。
  • Reduced是上一层减少的FLOPs。
  • Rest是上一层遗留下的计算次数。它们在传递给agent的时候会进行归一化。

在这个模型当中,作者考虑将上一层对本层的压缩影响考虑在内,而对本层量化主要考虑剪枝方式。通过减小nxc来获得压缩系数。

2、动作空间

在这里压缩系数是动作量,其对精度影响很大。如果采用离散动作空间,那么就需要划分很细,这样造成动作数量爆炸。所以作者采用了连续动作空间,动作量在(0, 1]之间。

3、DDPG

Agent获得每层的状态s_t,产生压缩率a_t。直到所有层压缩完成,压缩的网络在测试集上进行评估,将精度作为奖励返回到agent,agent更新自身参数,进行下一次强化学习,直到能达到一个合理的压缩系数和测试精度。作者在这里并没有对压缩后的网络再进行重新训练来微调参数,这样可以大大减少计算量,否则会需要花费很长时间来完成压缩。

三、蒸馏N2N learning

在知识蒸馏中,教师网络是通过人工修改来获得学生网络的,这篇文章让强化学习来做这一切。解除了人工干预,同时能够很好的遍历学生网络空间,照到更合适的学生网络。

强化学习策略输出的依然是一个压缩动作,可能是对卷积网络中某一个通道的删除。降低的网络就是学生网络空间中的一个元素。即将一个网络状态从s迁移到了s’。

但是这样的动作行为非常多,是和网络模型大小呈指数关系的。

作者为了降低行为空间的搜索量,采用了两步策略。

第一步用于移除教师网络中的层。
第二个在移除后的网络中进行层压缩。

生成的学生网络用知识蒸馏的方法进行训练。然后测试其精度作为奖励反馈给策略网络,策略网络利用随机策略梯度方法进行参数更新。

image.png

作者用马尔科夫模型来模拟寻找压缩模型过程,这个过程用五元组参量表示:
image.png

1、状态

状态S是包好所有压缩的可能子网络,比如从老师网络中移除一个filter,就是一个新的子网络,即一个新状态s。

2、动作

主要包含两方面:一个是移除网络层,另外一个是对移除后的网络进行层压缩。

3、转移函数

一个网络在下一个行动下生成一个新的网络。


image.png

4、discount函数

不同步下的行为产生的奖励对总奖励贡献不同的表示量。

5、奖励

image.png

C是压缩系数,A是子网络测试精度,用其和教师网络的测试精度比值来作为奖励。这里表示压缩越大或者精度越高奖励就越大,但是这里对压缩系数有一些限制,不能无限制压缩。


1、马尔科夫决策过程

因为使用的是强化学习的方法,所以我们首先定义一个马尔科夫的过程。

M = {S, A, T, r, γ}

其中

  • S是状态,即网络结构。
  • A是对网络结构的操作,这里的A有两个操作,一种是移除某一层,一种是缩小某一层。
  • T代表转移操作,即原状态s经过操作a后进入下一个状态s1。
  • r是回报函数,即做了操作后会得到什么收益。
  • γ是折扣系数,即多步以后的收益反映到当前步骤所乘以的系数,在这个问题中,折扣系数是1。

2、网络压缩

有了马尔科夫过程之后,我们先对网络做层次删减操作,然后再对删减后的网络做层次缩小操作。得到的新的小模型进行知识蒸馏学习,拿到得到的结果去计算回报。
整个过程如下图所示:

image.png

那么,强化学习中是如何控制层次删减和层次缩小的呢?在论文中,这两个的控制使用的是两个LSTM来实现。如下图
image.png

其中,左图用来控制某一层是不是该删除,输入是层次信息,包括层次类型,kernel size,stride,padding,output大小等,如果是残差连接的话,还有残差的起始层和结束层。

这个网络有L步,L就是teacher模型的层次数目。每一步的输入就是从低到高的某一层,然后得到这个LSTM的所有输出后,我们把要删减的层次删掉,就得到了新的模型。

得到了网络删减后的模型后,就用右图去给这个新模型进行瘦身。这个网络的输入仍然是每一层的具体信息,所不同的是,输入还会拼接上上一步的输出。输出则是一个(0, 1)区间内的数字,表示要缩小的比例。同样的,得到了这个模型所有步的输出后,就可以对网络做缩小操作,又得到了新网络。

最后,在用知识蒸馏去学习这个新网络。得到的结果去计算回报函数。用回报函数去训练这两个LSTM控制器。需要注意的是,训练这两个LSTM的时候,只有最后一步的输出上是要去计算损失函数的。


image.png

3、回报函数

image.png

其中C表示的压缩率,如下图


image.png

A表示的模型的准确率。

为什么是C(2-C),这个函数在0-1之间是递增函数,表示压缩率越高越好。 而A/A(teacher)表示,student模型越接近teacher模型越好。

4、优化

优化的方法就是使用强化学习的策略梯度方法:


image.png

对于回报高的方法,我们希望它的概率大。

而为了加速训练,一般会采用一个小技巧,那就是减去一个之前结果的均值,让回报函数有正有负,正的是我们需要的,负的是我们不需要的。


image.png

5、知识蒸馏

这块其实很简单,就是让student模型去学习teacher模型的输出。在论文中,知识蒸馏的损失函数设置的是logits的均方差。


image.png

四、量化Hardware-Aware Automated Quantization with Mixed Precision(HAQ)

(一)、模型量化是什么?本文研究了什么问题?

模型量化是对深度神经网络(DNN)进行压缩和加速的一种广泛使用的技术。在许多实时机器学习应用(如自动驾驶)中,DNN 受到延迟、能量和模型大小的严格限制。为了提高硬件的效率,许多研究者都提出将权值和激活值量化到低精度。

image.png

随着混合精度硬件的出现,需要提出混合精度量化方法。在传统的量化方法中,DNN 所有层的权重和激活值使用固定数目的 bit 位,如图 1(a) 所示,固定精度量化为每层的权重和激活值都分配了 8bit。但是由于不同层有不同的冗余,并且在硬件上表现不同,为每层分配相同数量的 bit 位并不是最优的,因此需要对不同层使用混合精度量化,如图 1(b) 所示。

在 18 年 8 月 NVIDIA 推出了支持 1-bit,4-bit,8-bit 和 16-bit 精度操作的图灵 GPU 架构;在 18 年 9 月,苹果发布了支持神经网络推理混合精度的 A12Bionic 芯片;除了工业界,学术界也开始研究 bit 级的灵活硬件设计:BISMO 提出了位串行乘法器以支持 1-8bits 的乘法;BitFusion 支持 2,4,8,16bits 的乘法。随着这些混合精度硬件及算法的出现,急需提出混合精度量化方法!

如何确定在不同的硬件加速器上每一层的权重和激活的位宽。这些硬件加速器的出现是为了进一步提高计算效率,然而也为寻找每层的最佳位宽提出了巨大挑战:它需要领域专家 (既有机器学习知识又有硬件体系结构知识) 用基于规则的启发式算法巧妙地探索巨大的设计空间,并且在延迟、能耗、模型尺寸等方面进行权衡。例如:我们应该在提取低层特征的第一层和计算最终输出的最后一层保留更多的 bit;

此外,我们应该在卷积层使用更多的 bit,而不是在全连接层,因为从经验上来说,卷积层更敏感。随着神经网络变得越来越深,搜索空间呈指数级增长,这使得依赖手工制作的量化策略变得不可行。因此,这些基于规则的量化策略通常是次优的,不能从一个模型推广到另一个模型。如果我们有 M 个不同的神经网络模型,每一个模型有 N 层,在 H 个不同的硬件平台上,可能的解决方案就有 O(H×M×8^(2N))。如何简化这一步骤,自动化探索不同硬件加速器上每一层权重和激活的位宽,是迫切需要的解决的!

如何在硬件上优化给定模型的延迟和能量消耗。一种被广泛采用的方法是依赖一些 Proxy 信号 (FLOPs,内存引用的数量)。然而,由于不同的硬件表现非常不同,模型在硬件上的性能并不总是能够由这些 Proxy 信号准确地反映出来。因此,直接涉及硬件体系结构,在硬件上优化给定模型的延迟和能量消耗就显得尤为重要!

(二)、HAQ 用什么方法解决了上述问题?

作者提出了硬件感知自动量化(HAQ)框架(如图 2 所示),将量化任务建模为强化学习问题,自动搜索量化策略。通过深度确定性策略梯度 (DDPG) 来监督 agent,依赖于 actor-critic 模型,在给定计算资源量(延迟,能耗和模型尺寸)的情况下,agent 依赖硬件加速器的直接反馈(延迟和能耗),而不是依赖 proxy 信号(如 FLOPs 和模型尺寸),按层处理给定的神经网络,接收相应信息作为 observation,给出每个层的 action:即每层的 bits,并将准确率作为 reward 反馈给 critic,从而确定最佳位宽分配策略。

Agent 将按层处理给定的神经网络。对于每个层,agent 接收层配置和统计信息作为 observation,用一个十维特征向量 O_k 表示

  • 如果第 k 层是卷积层,那么 O_k 是:


    image.png
  • 如果第 k 层是全连接层,那么 O_k 是:


    image.png

    对于 O_k 中的每个维度,作者都归一化到了 [0,1],使得它们的规模相同。

根据 Observation,作者使用连续 action 空间来确定位宽,在第 k 个 time step,作者会将 a_k 四舍五入为离散的位宽值 b_k:


image.png

其中 bmin 和 bmax 表示最小和最大位宽(在实验中,作者将 bmin 设置为 2,bmax 设置为 8)。

在实际应用中,具有有限的计算预算(即延迟,能量和模型大小)。作者希望找到具有约束条件的最佳性能的量化策略,因此鼓励 agent 通过限制 action 空间来满足计算资源预算。在 agent 向所有层提供 action{a_k} 之后,测量量化模型将使用的资源量。直接从硬件加速器获得反馈,如果当前量化策略超出资源预算(延迟,能耗或模型尺寸),则将按顺序减小每个层的位宽,直到最终满足约束。

对 agent 的直接反馈可以是 proxy 信号 (Flops 或模型尺寸),然而由于 proxy 信号是间接的,它们不等同于硬件上的性能(延迟,能耗)。同时局部缓存,内核调用次数,内存带宽等信息也都很重要。而 Proxy 反馈无法对这些硬件功能进行建模以找到特定的量化策略。因此,作者使用来自硬件加速器的直接延迟和能量反馈作为资源约束,这使得 agent 能够根据不同层之间的细微差别来确定位宽分配策略。在量化所有层之后,作者将量化模型微调一个 epoch,并在短期再训练之后将验证准确率作为 reward 提供给 agent。reward 的计算如下所示:


image.png

其中,acc_quant 是微调后量化模型的准确率,acc_origin 是 full-precision 模型在训练集上的 top-1 准确率,λ是缩放因子,在实验中作者设置为 0.1。


  • 描述成一个寻找特定constraint (Latency/Energy)下的最小accuracy degradation的网络问题
  • weight和activation使用不同的量化bit-width


    image.png
  • 使用RL


    image.png

image.png

image.png

image.png
©著作权归作者所有,转载或内容合作请联系作者
  • 序言:七十年代末,一起剥皮案震惊了整个滨河市,随后出现的几起案子,更是在滨河造成了极大的恐慌,老刑警刘岩,带你破解...
    沈念sama阅读 158,736评论 4 362
  • 序言:滨河连续发生了三起死亡事件,死亡现场离奇诡异,居然都是意外死亡,警方通过查阅死者的电脑和手机,发现死者居然都...
    沈念sama阅读 67,167评论 1 291
  • 文/潘晓璐 我一进店门,熙熙楼的掌柜王于贵愁眉苦脸地迎上来,“玉大人,你说我怎么就摊上这事。” “怎么了?”我有些...
    开封第一讲书人阅读 108,442评论 0 243
  • 文/不坏的土叔 我叫张陵,是天一观的道长。 经常有香客问我,道长,这世上最难降的妖魔是什么? 我笑而不...
    开封第一讲书人阅读 43,902评论 0 204
  • 正文 为了忘掉前任,我火速办了婚礼,结果婚礼上,老公的妹妹穿的比我还像新娘。我一直安慰自己,他们只是感情好,可当我...
    茶点故事阅读 52,302评论 3 287
  • 文/花漫 我一把揭开白布。 她就那样静静地躺着,像睡着了一般。 火红的嫁衣衬着肌肤如雪。 梳的纹丝不乱的头发上,一...
    开封第一讲书人阅读 40,573评论 1 216
  • 那天,我揣着相机与录音,去河边找鬼。 笑死,一个胖子当着我的面吹牛,可吹牛的内容都是我干的。 我是一名探鬼主播,决...
    沈念sama阅读 31,847评论 2 312
  • 文/苍兰香墨 我猛地睁开眼,长吁一口气:“原来是场噩梦啊……” “哼!你这毒妇竟也来了?” 一声冷哼从身侧响起,我...
    开封第一讲书人阅读 30,562评论 0 197
  • 序言:老挝万荣一对情侣失踪,失踪者是张志新(化名)和其女友刘颖,没想到半个月后,有当地人在树林里发现了一具尸体,经...
    沈念sama阅读 34,260评论 1 241
  • 正文 独居荒郊野岭守林人离奇死亡,尸身上长有42处带血的脓包…… 初始之章·张勋 以下内容为张勋视角 年9月15日...
    茶点故事阅读 30,531评论 2 245
  • 正文 我和宋清朗相恋三年,在试婚纱的时候发现自己被绿了。 大学时的朋友给我发了我未婚夫和他白月光在一起吃饭的照片。...
    茶点故事阅读 32,021评论 1 258
  • 序言:一个原本活蹦乱跳的男人离奇死亡,死状恐怖,灵堂内的尸体忽然破棺而出,到底是诈尸还是另有隐情,我是刑警宁泽,带...
    沈念sama阅读 28,367评论 2 253
  • 正文 年R本政府宣布,位于F岛的核电站,受9级特大地震影响,放射性物质发生泄漏。R本人自食恶果不足惜,却给世界环境...
    茶点故事阅读 33,016评论 3 235
  • 文/蒙蒙 一、第九天 我趴在偏房一处隐蔽的房顶上张望。 院中可真热闹,春花似锦、人声如沸。这庄子的主人今日做“春日...
    开封第一讲书人阅读 26,068评论 0 8
  • 文/苍兰香墨 我抬头看了看天上的太阳。三九已至,却和暖如春,着一层夹袄步出监牢的瞬间,已是汗流浃背。 一阵脚步声响...
    开封第一讲书人阅读 26,827评论 0 194
  • 我被黑心中介骗来泰国打工, 没想到刚下飞机就差点儿被人妖公主榨干…… 1. 我叫王不留,地道东北人。 一个月前我还...
    沈念sama阅读 35,610评论 2 274
  • 正文 我出身青楼,却偏偏与公主长得像,于是被迫代替她去往敌国和亲。 传闻我的和亲对象是个残疾皇子,可洞房花烛夜当晚...
    茶点故事阅读 35,514评论 2 269

推荐阅读更多精彩内容