pytorch模型剪枝学习笔记

pytorch代码仓库

pytorch在19年11月份的时候合入了这部分剪枝的代码。pytorch提供一些直接可用的api,用户只需要传入需要剪枝的module实例和需要剪枝的参数名字,系统自动帮助完成剪枝操作,看起来接口挺简单。比如 def random_structured(module, name, amount, dim)

pytorch支持的几种类型的剪枝策略:

详细分析

  • pytorch提供了一个剪枝的抽象基类‘‘class BasePruningMethod(ABC)’,所有剪枝策略都需要继承该基类,并重载部分函数就可以了

  • 一般情况下需要重载init和compute_mask方法,call, apply_mask, apply, prune和remove不需要重载,例如官方提供的RandomUnstructured剪枝方法

    file

  • 基类实现的6个方法:


    file
  • 剪枝的API接口,可以看到支持用户自定义的剪枝mask,接口为custom_from_mask


    file
  • API的实现,使用classmethod的方法,剪枝策略的实例化在框架内部完成,不需要用户实例化

  • 剪枝的大只过程:

    1. 根据用户选择的剪枝API生成对应的策略实例,此时会判断需要做剪枝操作的module上是否已经挂有前向回调函数,没有则生成新的,有了就在老的上面添加,并且生成PruningContainer。从这里可以看出,对于同一个module使用多个剪枝策略时,pytorch通过PruningContainer来对剪枝策略进行管理。PruningContainer本身也是继承自BasePruningMethod。同时设置前向计算的回调,便于后续训练时调用。
    2. 接着根据用户输入的module和name,找到对应的参数tensor。如果是第一次剪枝,那么需要生成_orig结尾的tensor,然后删除原始的module上的tensor。如name为bias,那么生成bias_orig存起来,然后删除module.bias属性。
    3. 获取defaultmask,然后调用method.computemask生成当前策略的mask值。生成的mask会被存在特定的缓存module.register_buffer(name + "_mask", mask)。这里的compute_mask可能是两种情况:如果只有一个策略,那么调用的时候对应剪枝策略的compute_mask方法,如果一个module有多个剪枝策略组合,那么调用的应该是PruningContainer的compute_mask


      file
    4. 执行剪枝,保存剪枝结果到module的属性,注册训练时的剪枝回调函数,剪枝完成。新的mask应用在orig的tensor上面生成新的tensor保存的对应的name属性


      file
  • remove接口
    pytorch还提供各类一个remove接口,目的是把之前的剪枝结果持久化,具体操作就是删除之前生成的跟剪枝相关的缓存或者是回调hook接口,设置被剪枝的name参数(如bias)为最后一次训练的值。

    file

  • 自己写一个剪枝策略接口也是可以的:


    file
    1. 先写一个剪枝策略类继承BasePruningMethod
    2. 然后重载基类的compute_mask方法,写自己的计算mask方法

官方完整教程在这里

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