https://github.com/youzhonghui/gate-decorator-pruning
1.utils.py
class dotdict(dict): """dot.notation access to dictionary attributes""" __getattr__ = dict.get __setattr__ = dict.__setitem__ __delattr__ = dict.__delitem__
继承dict字典，实际上还是dict
2.loader/__init__.py
import torchvision from torch.utils.data import DataLoader from torchvision import transforms from PIL import Image from config import cfg from loader.cifar10 import get_cifar10 from loader.cifar100 import get_cifar100 from loader.imagenet import get_imagenet def get_loader(): pair = { # 根据设置的参数的数据集名字来选择使用哪个数据集 'cifar10': get_cifar10, 'cifar100': get_cifar100, 'imagenet': get_imagenet } return pair[cfg.data.type]()
选择使用那个数据集，对应的config设置为：
from config import parse_from_dict parse_from_dict({ ... "data": { "type": "cifar10", #这个即使用的数据集名字 "shuffle": True, "batch_size": 128, "test_batch_size": 128, "num_workers": 4 ...
3.models/__init__.py
import torch from config import cfg def get_vgg16_for_cifar(): from models.cifar.vgg import VGG return VGG('VGG16', cfg.model.num_class) def get_resnet50_for_imagenet(): from models.imagenet.resnet50 import Resnet50 return Resnet50(cfg.model.num_class) def get_resnet56(): from models.cifar.resnet56 import resnet56 return resnet56(cfg.model.num_class) def get_model(): pair = {# 根据设置的参数的模型名字来选择使用哪个模型 'cifar.vgg16': get_vgg16_for_cifar, 'resnet50': get_resnet50_for_imagenet, 'cifar.resnet56': get_resnet56 } model = pair[cfg.model.name]() if cfg.base.checkpoint_path != '': #是否有训练好的预训练模型 print('restore checkpoint: ' + cfg.base.checkpoint_path) model.load_state_dict(torch.load(cfg.base.checkpoint_path, map_location='cpu' if not cfg.base.cuda else 'cuda')) if cfg.base.cuda: #单个GPU model = model.cuda() if cfg.base.multi_gpus: #多个GPU model = torch.nn.DataParallel(model) return model
选择使用哪个模型进行分类，并设置是使用cpu还是GPU，有预训练模型就加载预训练模型
对应的config设置为：
from config import parse_from_dict parse_from_dict({ ... "model": { "name": "cifar.resnet56", "num_class": 10, "pretrained": False },
4.loss.py
import torch import torch.nn as nn import torch.nn.functional as F from torch.autograd import Variable import numpy as np from config import cfg def get_criterion(): pair = { #设置使用的损失函数 'softmax': nn.CrossEntropyLoss() } assert (cfg.loss.criterion in pair) criterion = pair[cfg.loss.criterion] return criterion
使用交叉熵损失函数
对应的config设置为：
from config import parse_from_dict parse_from_dict({ ... "loss": { "criterion": "softmax"
5. config.py
import argparse import json from utils import dotdict def make_as_dotdict(obj): #从dict改成dotdict格式 if type(obj) is dict: obj = dotdict(obj) for key in obj: if type(obj[key]) is dict: obj[key] = make_as_dotdict(obj[key]) return obj def parse(): print('Parsing config file...') parser = argparse.ArgumentParser(description="config") parser.add_argument( "--config", type=str, default="configs/base.json", help="Configuration file to use" ) cli_args = parser.parse_args() with open(cli_args.config) as fp: config = make_as_dotdict(json.loads(fp.read())) print(json.dumps(config, indent=4, sort_keys=True)) return config class Singleton(object): _instance = None def __new__(cls, *args, **kw): if not cls._instance: cls._instance = super(Singleton, cls).__new__(cls, *args, **kw) return cls._instance class Config(Singleton): def __init__(self): self._cfg = dotdict({}) try: self._cfg = parse() except: pass def __getattr__(self, name): if name == '_cfg': super().__setattr__(name) else: return self._cfg.__getattr__(name) def __setattr__(self, name, val): if name == '_cfg': super().__setattr__(name, val) else: self._cfg.__setattr__(name, val) def __delattr__(self, name): #del删除元素时调用 return self._cfg.__delitem__(name) def copy(self, new_config): self._cfg = make_as_dotdict(new_config) cfg = Config() def parse_from_dict(d): #将dict换成dotdict global cfg assert type(d) == dict cfg.copy(d)
设置参数
但是不太明白为什么要弄成dotdict格式
这个函数在后面进行prune和finetune的时候会调用来设置参数信息，如：
from config import parse_from_dict parse_from_dict({ "base": { "task_name": "resnet56_cifar10_ticktock", "cuda": True, "seed": 0, "checkpoint_path": "", "epoch": 0, "multi_gpus": True, "fp16": False }, "model": { "name": "cifar.resnet56", "num_class": 10, "pretrained": False }, "train": { "trainer": "normal", "max_epoch": 160, "optim": "sgd", "steplr": [ [80, 0.1], #step>=80时,学习率都设置为0.1 [120, 0.01], [160, 0.001] # 120<step<=160时将学习率设置为0.001 ], "weight_decay": 5e-4, "momentum": 0.9, "nesterov": False }, "data": { "type": "cifar10", "shuffle": True, "batch_size": 128, "test_batch_size": 128, "num_workers": 4 }, "loss": { "criterion": "softmax" }, "gbn": { "sparse_lambda": 1e-3, "flops_eta": 0, "lr_min": 1e-3, "lr_max": 1e-2, "tock_epoch": 10, "T": 10, "p": 0.002 } }) from config import cfg
6.
trainer/__init__.py
from trainer.normal import NormalTrainer
from config import cfg
def get_trainer():
pair = {
'normal': NormalTrainer
}
assert (cfg.train.trainer in pair)
return pair[cfg.train.trainer]()
设置使用的训练train()、测试test()函数所在位置
扩展：
#coding:utf-8 import torch if __name__ == '__main__': a = torch.FloatTensor([[3, 14, 15, 13], [5,4,15,7]]).t() b = torch.FloatTensor([[3, 3, 3, 3], [5,5,5,5]]).t() correct = a.eq(b) print(correct) print(correct[:1]) print(correct[:1].view(-1)) print(correct[:1].view(-1).float()) correct_k = correct[:1].view(-1).float().sum(0, keepdim=True) print(correct_k)
返回：
tensor([[ True, True], [False, False], [False, False], [False, False]]) tensor([[True, True]]) tensor([True, True]) tensor([1., 1.]) tensor([2.])
trainer/normal.py：
from time import time import torch import torch.nn as nn from torch.autograd import Variable import torch.nn.functional as F from tqdm import tqdm import numpy as np from config import cfg FINISH_SIGNAL = 'finish' def accuracy(output, target, topk=(1,)): #计算分类的准确度 """Computes the accuracy over the k top predictions for the specified values of k""" with torch.no_grad(): maxk = max(topk) #看是top-1还是top-5 batch_size = target.size(0) _, pred = output.topk(maxk, 1, True, True) # 从输出中得到前maxk个大的预测结果的索引值,大小为(batch_size, maxk) pred = pred.t() # 转置成(maxk, batch_size) # target从(batch_size, 1) -> (1, batch_size) -> (maxk, batch_size) # 然后与pred对比看是否相等，每个batch_size最多只有一个相等，所以correct中true的个数最大值为batch_size # correct为(maxk, batch_size)，值为 correct = pred.eq(target.view(1, -1).expand_as(pred)) res = [] for k in topk: correct_k = correct[:k].view(-1).float().sum(0, keepdim=True) res.append(correct_k.mul_(100.0 / batch_size)) # 得到百分比的正确率 return res class NormalTrainer(): def __init__(self): self.use_cuda = cfg.base.cuda def test(self, pack, topk=(1,)): # 测试 pack.net.eval() loss_acc, correct, total = 0.0, 0.0, 0.0 hub = [[] for i in range(len(topk))] for data, target in pack.test_loader: if self.use_cuda: data, target = data.cuda(), target.cuda() with torch.no_grad(): #不后向传播 output = pack.net(data) loss_acc += pack.criterion(output, target).data.item() #计算损失 acc = accuracy(output, target, topk) # 准确率 for acc_idx, score in enumerate(acc): hub[acc_idx].append(score[0].item()) loss_acc /= len(pack.test_loader) # 最后得到的平均损失 info = { 'test_loss': loss_acc } for acc_idx, k in enumerate(topk): info['acc@%d' % k] = np.mean(hub[acc_idx]) #top-1,top-5等准确率 return info def train(self, pack, loss_hook=None, iter_hook=None, update=True, mute=False, acc_step=1): #训练，mute即是否打印info pack.net.train() loss_acc, correct_acc, total = 0.0, 0.0, 0.0 begin = time() pack.optimizer.zero_grad() with tqdm(total=len(pack.train_loader), disable=mute) as pbar: total_iter = len(pack.train_loader) #总迭代次数 for cur_iter, (data, label) in enumerate(pack.train_loader): if iter_hook is not None: signal = iter_hook(cur_iter, total_iter) if signal == FINISH_SIGNAL: #结束标志 break if self.use_cuda: data, label = data.cuda(), label.cuda() data = Variable(data, requires_grad=False) label = Variable(label) logits = pack.net(data) loss = pack.criterion(logits, label) if loss_hook is not None: additional = loss_hook(data, label, logits) loss += additional loss = loss / acc_step loss.backward() if (cur_iter + 1) % acc_step == 0: if update: pack.optimizer.step() pack.optimizer.zero_grad() loss_acc += loss.item() pbar.update(1) info = { 'train_loss': loss_acc / len(pack.train_loader), 'epoch_time': time() - begin } return info
一个train()即跑完一次所有数据就结束了，即enumerate(pack.train_loader)完就完了
7.main.py
""" * Copyright (C) 2019 Zhonghui You * If you are using this code in your research, please cite the paper: * Gate Decorator: Global Filter Pruning Method for Accelerating Deep Convolutional Neural Networks, in NeurIPS 2019. """ import torch import torch.nn as nn import torch.optim as optim import numpy as np import random import math from loader import get_loader from models import get_model from trainer import get_trainer from loss import get_criterion from utils import dotdict from config import cfg from logger import logger def _sgdr(epoch): lr_min, lr_max = cfg.train.sgdr.lr_min, cfg.train.sgdr.lr_max restart_period = cfg.train.sgdr.restart_period _epoch = epoch - cfg.train.sgdr.warm_up while _epoch/restart_period > 1.: _epoch = _epoch - restart_period restart_period = restart_period * 2. radians = math.pi*(_epoch/restart_period) return lr_min + (lr_max - lr_min) * 0.5*(1.0 + math.cos(radians)) def _step_lr(epoch): v = 0.0 for max_e, lr_v in cfg.train.steplr: #max_e是到这个step的学习率都是lr_v v = lr_v if epoch <= max_e: break return v def get_lr_func(): if cfg.train.steplr is not None: return _step_lr elif cfg.train.sgdr is not None: return _sgdr else: assert False def adjust_learning_rate(epoch, pack): #设置使用的优化器，并设置学习率调节函数，以及更新学习率 if pack.optimizer is None: if cfg.train.optim == 'sgd' or cfg.train.optim is None: pack.optimizer = optim.SGD( pack.net.parameters(), lr=1, momentum=cfg.train.momentum, weight_decay=cfg.train.weight_decay, nesterov=cfg.train.nesterov ) else: print('WRONG OPTIM SETTING!') assert False pack.lr_scheduler = optim.lr_scheduler.LambdaLR(pack.optimizer, get_lr_func()) pack.lr_scheduler.step(epoch) return pack.lr_scheduler.get_lr() def recover_pack(): train_loader, test_loader = get_loader() pack = dotdict({ 'net': get_model(), 'train_loader': train_loader, 'test_loader': test_loader, 'trainer': get_trainer(), 'criterion': get_criterion(), 'optimizer': None, 'lr_scheduler': None }) adjust_learning_rate(cfg.base.epoch, pack) return pack def set_seeds(): #用来保证代码中随机数每次都一样 torch.manual_seed(cfg.base.seed) if cfg.base.cuda: torch.cuda.manual_seed_all(cfg.base.seed) torch.backends.cudnn.deterministic = True if cfg.base.fp16: torch.backends.cudnn.enabled = True # torch.backends.cudnn.benchmark = True np.random.seed(cfg.base.seed) random.seed(cfg.base.seed) def main(): set_seeds() #设置中设置的"seed": 0,就是用在这的 pack = recover_pack() #设置各个参数和使用的模型、数据等 for epoch in range(cfg.base.epoch + 1, cfg.train.max_epoch + 1): lr = adjust_learning_rate(epoch, pack) # 更新lr info = pack.trainer.train(pack) #训练模型，得到损失和准确率等信息 info.update(pack.trainer.test(pack)) #加入测试时的损失和准确率等信息 info.update({'LR': lr}) #记录此时的lr print(epoch, info) logger.save_record(epoch, info) #写入日志 if epoch % cfg.base.model_saving_interval == 0: logger.save_network(epoch, pack.net) # 保存网络 if __name__ == '__main__': main()
8.logger.py
import torch from config import cfg import os import json import numpy as np class MetricsRecorder(): def __init__(self): self.rec = {} def add(self, pairs): for key, val in pairs.items(): if key not in self.rec: self.rec[key] = [] self.rec[key].append(val) def mean(self): r = {} for key, val in self.rec.items(): r[key] = np.mean(val) return r class Logger(): def __init__(self): self.base_path = './logs/' + cfg.base.task_name self.logfile = self.base_path + '/log.json' self.cfgfile = self.base_path + '/cfg.json' if not os.path.isdir(self.base_path): os.makedirs(self.base_path, exist_ok=True) with open(self.logfile, 'w') as fp: json.dump({}, fp) #初始化时日志信息为空 with open(self.cfgfile, 'w') as fp: json.dump(cfg, fp) #初始化时配置信息即config信息 def save_record(self, epoch, record): #保存运行过程中训练和测试的损失和准确率等信息，以当前的epoch为索引 with open(self.logfile) as fp: log = json.load(fp) log[str(epoch)] = record with open(self.logfile, 'w') as fp: json.dump(log, fp) def save_network(self, epoch, network): saving_path = self.base_path + '/ckp.%d.torch' % epoch print('saving model ...') if type(network) is torch.nn.DataParallel: torch.save(network.module.state_dict(), saving_path) else: torch.save(network.state_dict(), saving_path) cfg.base.epoch = epoch cfg.base.checkpoint_path = saving_path with open(self.cfgfile, 'w') as fp: # 保存新的配置信息 json.dump(cfg, fp) logger = None if logger is None: logger = Logger()
相应的文件将会根据任务名字，即设置：
from config import parse_from_dict parse_from_dict({ "base": { "task_name": "resnet18", #任务名字
在./logs文件夹下创建同名文件夹存储log.json和cfg.json文件，save_record()就是将中间信息保存在这，调用save_network()也会将模型保存在该文件夹中
user@jiayuan:/opt/.../gate-decorator-pruning/logs/resnet18$ ls
cfg.json log.json
接下来就是prune和finetune了,重要
9.prune/utils.py
#coding:utf-8 import os if __name__ == '__main__': print(os.devnull) #/dev/null
代码：
import torch import torch.nn as nn import os, contextlib from thop import profile def analyse_model(net, inputs): # silence with open(os.devnull, 'w') as devnull: #os.devnull对于Linux为/dev/null with contextlib.redirect_stdout(devnull):#标准输出已经重定向到了 /dev/null flops, params = profile(net, (inputs, )) #估算PyTorch模型的FLOPs模块 return flops, params def finetune(pack, lr_min, lr_max, T, mute=False): #T即finetune_epoch，即40轮迭代 logs = [] epoch = 0 def iter_hook(curr_iter, total_iter): #作为train的iter_hook参数传入 total = T * total_iter #total_iter即dataloader中有多少批batch_size，所以整个finetune跑total个batch_size half = total / 2 itered = epoch * total_iter + curr_iter #curr_iter即一个epoch中，数据跑到了第curr_iter个batch_size，现在的总batch_size数为itered if itered < half: #当小于一半时，学习率这么算 _iter = epoch * total_iter + curr_iter _lr = (1- _iter / half) * lr_min + (_iter / half) * lr_max else: # 当大于或等于一半时，学习率这么算，这两个的差别就是lr_max和lr_min的前后位置不同，大概意思是相同的 _iter = (epoch - T/2) * total_iter + curr_iter _lr = (1- _iter / half) * lr_max + (_iter / half) * lr_min for g in pack.optimizer.param_groups: g['lr'] = max(_lr, 0) g['momentum'] = 0.0 for i in range(T): #训练40个epoch info = pack.trainer.train(pack, iter_hook = iter_hook) info.update(pack.trainer.test(pack)) info.update({'LR': pack.optimizer.param_groups[0]['lr']}) epoch += 1 if not mute: #是否打印损失和精确度等信息 print(info) logs.append(info) return logs
这里的微调操作其实就是论文中的：
它跟tick-tock中的tock的差别在于tock中使用的还是GBN，且训练次数比较少，一半就10次；而finetune操作是在整个模型都prune后的小模型中训练，GBN都换回BN，且训练次数也比较多
10.prune/universal.py
扩展
1)uuid库：
UUID: 通用唯一标识符 ( Universally Unique Identifier ), 对于所有的UUID它可以保证在空间和时间上的唯一性. 它是通过MAC地址, 时间戳, 命名空间, 随机数, 伪随机数来保证生成ID的唯一性, 有着固定的大小( 128 bit ). 它的唯一性和一致性特点使得可以无需注册过程就能够产生一个新的UUID. UUID可以被用作多种用途, 既可以用来短时间内标记一个对象, 也可以可靠的辨别网络中的持久性对象.
为什么要使用UUID?
很多应用场景需要一个id, 但是又不要求这个id 有具体的意义, 仅仅用来标识一个对象. 常见的例子有数据库表的id 字段. 另一个例子是前端的各种UI库, 因为它们通常需要动态创建各种UI元素, 这些元素需要唯一的id , 这时候就需要使用UUID了.
#coding:utf-8 import uuid if __name__ == '__main__': print(uuid.uuid1()) #7b24099a-27ae-11ea-b076-00e04c6841ff
其实这个库主要是用于像resnet这样的网络中有侧枝shortcut的情况，是分组使用的，即同一个Group的ID是相同的。像VGG这样的网络每个GBN层的ID是不同的
2)nn.Parameter
#coding:utf-8 import torch.nn as nn import torch if __name__ == '__main__': g = nn.Parameter(torch.ones(1, 3, 1, 1), requires_grad=True) print(g)
返回：
Parameter containing: tensor([[[[1.]], [[1.]], [[1.]]]], requires_grad=True)
使用nn.Parameter的目的是将一个不可训练的类型Tensor转换成可以训练的类型parameter并将这个parameter绑定到这个module里面(net.parameter()中就有这个绑定的parameter，所以在参数优化的时候可以进行优化的)，所以经过类型转换这个值就变成了模型的一部分，成为了模型中根据训练可以改动的参数了
模型中的bias和weight都是nn.Parameter，可用于训练，并实现优化；Variable则是作为模型的输入
buffers()返回一个模块缓冲区的迭代器,其保存的是模型中每次前向传播需用到上一次前向传播的结果，作为持久状态的值，如BatchNorm2d()中使用的均值和方差值,其随着BatchNorm2d()中参数的变化而变化
3）
所以GatedBatchNorm2d代码中初始化中有设置参数：
def extract_from_bn(self):
# freeze bn weight
with torch.no_grad():
self.bn.bias.set_(torch.clamp(self.bn.bias / self.bn.weight, -10, 10)) #将self.bn.bias / self.bn.weight的值保持在[-10, 10],小于-10的即改为-10，大于10的即改为10
self.g.set_(self.g * self.bn.weight.view(1, -1, 1, 1))
self.bn.weight.set_(torch.ones_like(self.bn.weight)) #torch.ones_like(input)相当于torch.ones(input.size())
self.bn.weight.requires_grad = False
Φ就是g，β就是bn.bias，γ就是self.bn.weight
prune后，得到应该截取掉的filter，变回来的代码：
def melt(self): with torch.no_grad(): mask = self.bn_mask.view(-1) #转成列表， mask中有channels个值，值为0说明该channel被prune了 replacer = nn.BatchNorm2d(int(self.bn_mask.sum())).to(self.bn.weight.device) replacer.running_var.set_(self.bn.running_var[mask != 0]) #BatchNorm2d中的方差 replacer.running_mean.set_(self.bn.running_mean[mask != 0]) #BatchNorm2d中的均值 replacer.weight.set_((self.bn.weight * self.g.view(-1))[mask != 0]) replacer.bias.set_((self.bn.bias * self.g.view(-1))[mask != 0]) return replacer
整个代码：
import torch import torch.nn as nn import numpy as np import uuid OBSERVE_TIMES = 5 FINISH_SIGNAL = 'finish' class Meltable(nn.Module): def __init__(self): super(Meltable, self).__init__() @classmethod def melt_all(cls, net): def _melt(modules): keys = modules.keys() for k in keys: if len(modules[k]._modules) > 0: _melt(modules[k]._modules) if isinstance(modules[k], Meltable): modules[k] = modules[k].melt() _melt(net._modules) @classmethod def observe(cls, pack, lr): tmp = pack.train_loader if pack.tick_trainset is not None: pack.train_loader = pack.tick_trainset #数据集 for m in pack.net.modules(): if isinstance(m, nn.BatchNorm2d): m.weight.data.abs_().add_(1e-3) def replace_relu(modules): #将relu函数换成LeakyReLU函数 keys = modules.keys() for k in keys: if len(modules[k]._modules) > 0: replace_relu(modules[k]._modules) if isinstance(modules[k], nn.ReLU): modules[k] = nn.LeakyReLU(inplace=True) replace_relu(pack.net._modules) count = 0 def _freeze_bn(curr_iter, total_iter): for m in pack.net.modules(): if isinstance(m, nn.BatchNorm2d): m.eval() nonlocal count count += 1 if count == OBSERVE_TIMES: return FINISH_SIGNAL info = pack.trainer.train(pack, iter_hook=_freeze_bn, update=False, mute=True) #优化器不优化了 def recover_relu(modules): #将LeakyReLU函数换成relu函数 keys = modules.keys() for k in keys: if len(modules[k]._modules) > 0: recover_relu(modules[k]._modules) if isinstance(modules[k], nn.LeakyReLU): modules[k] = nn.ReLU(inplace=True) recover_relu(pack.net._modules) for m in pack.net.modules(): if isinstance(m, nn.BatchNorm2d): m.weight.data.abs_().add_(-1e-3) # 变回来 pack.train_loader = tmp class GatedBatchNorm2d(Meltable): def __init__(self, bn, minimal_ratio = 0.1): super(GatedBatchNorm2d, self).__init__() assert isinstance(bn, nn.BatchNorm2d) self.bn = bn self.group_id = uuid.uuid1() self.channel_size = bn.weight.shape[0] self.minimal_filter = max(1, int(self.channel_size * minimal_ratio)) #最小的通道数 self.device = bn.weight.device self._hook = None self.g = nn.Parameter(torch.ones(1, self.channel_size, 1, 1).to(self.device), requires_grad=True)#一个可以用于训练的参数 # 这样后就会生成三个参数self.areaself.scoreself.bn_mask self.register_buffer('area', torch.zeros(1).to(self.device)) #即nn.Module.register_buffer,保存一些前向传播会用到的上一次前向传播的结果 self.register_buffer('score', torch.zeros(1, self.channel_size, 1, 1).to(self.device)) self.register_buffer('bn_mask', torch.ones(1, self.channel_size, 1, 1).to(self.device)) #bn_mask就是用来记录该bn的channels层是否被prune了，为0则被prune了，初始化为1 self.extract_from_bn() #将本身的bn的weight、bias和g三个参数重新设置一下 def set_groupid(self, new_id): self.group_id = new_id def extra_repr(self): #即prune后channel数从channel_size变为了bn_mask.sum() return '%d -> %d | ID: %s' % (self.channel_size, int(self.bn_mask.sum()), self.group_id) def extract_from_bn(self): # freeze bn weight with torch.no_grad(): self.bn.bias.set_(torch.clamp(self.bn.bias / self.bn.weight, -10, 10)) #将self.bn.bias / self.bn.weight的值保持在[-10, 10],小于-10的即改为-10，大于10的即改为10 self.g.set_(self.g * self.bn.weight.view(1, -1, 1, 1)) self.bn.weight.set_(torch.ones_like(self.bn.weight)) #torch.ones_like(input)相当于torch.ones(input.size()) self.bn.weight.requires_grad = False def reset_score(self): self.score.zero_() def cal_score(self, grad): # used for hook self.score += (grad * self.g).abs() #论文中公式6的计算，计算分数，即变成prune后的网络和以前网络的损失差计算，得到此时设置的每个参数g的分数，分数越小直至0说明该g的channels能删除 def start_collecting_scores(self): if self._hook is not None: self._hook.remove() self._hook = self.g.register_hook(self.cal_score) #后向传播计算出关于这个参数g的gradient后将会调用cal_score计算此时的self.score分数值，排序使用 def stop_collecting_scores(self): if self._hook is not None: self._hook.remove() # 移除register_hook得到的hook self._hook = None def get_score(self, eta=0.0): # use self.bn_mask.sum() to calculate the number of input channel. eta should had been normed # 因为self.bn_mask中的值都是1，大小为torch.ones(1, self.channel_size, 1, 1)，所以sum()后的结果为self.channel_size flops_reg = eta * int(self.area[0]) * self.bn_mask.sum() return ((self.score - flops_reg) * self.bn_mask).view(-1) def forward(self, x): x = self.bn(x) * self.g # self.g就是用来排重要性的参数 self.area[0] = x.shape[-1] * x.shape[-2] #长*宽=面积area if self.bn_mask is not None: return x * self.bn_mask return x def melt(self): with torch.no_grad(): mask = self.bn_mask.view(-1) #得到当前prune后的channels数 replacer = nn.BatchNorm2d(int(self.bn_mask.sum())).to(self.bn.weight.device) replacer.running_var.set_(self.bn.running_var[mask != 0]) #BatchNorm2d中的方差 replacer.running_mean.set_(self.bn.running_mean[mask != 0]) #BatchNorm2d中的均值 replacer.weight.set_((self.bn.weight * self.g.view(-1))[mask != 0]) replacer.bias.set_((self.bn.bias * self.g.view(-1))[mask != 0]) return replacer @classmethod def transform(cls, net, minimal_ratio=0.1): r = [] def _inject(modules): keys = modules.keys() for k in keys: if len(modules[k]._modules) > 0: _inject(modules[k]._modules) if isinstance(modules[k], nn.BatchNorm2d): # 将模型中的nn.BatchNorm2d换成GatedBatchNorm2d，截取后的filter数量>= max(1, int(self.channel_size * minimal_ratio)) modules[k] = GatedBatchNorm2d(modules[k], minimal_ratio) r.append(modules[k]) _inject(net._modules) return r
4）
卷积层的prune：
class Conv2dObserver(Meltable): def __init__(self, conv): super(Conv2dObserver, self).__init__() assert isinstance(conv, nn.Conv2d) self.conv = conv self.in_mask = torch.zeros(conv.in_channels).to('cpu') self.out_mask = torch.zeros(conv.out_channels).to('cpu') self.f_hook = conv.register_forward_hook(self._forward_hook) #该层卷机前向传播是进行的操作 def extra_repr(self): return '(%d, %d) -> (%d, %d)' % (self.conv.in_channels, self.conv.out_channels, int((self.in_mask != 0).sum()), int((self.out_mask != 0).sum())) def _forward_hook(self, m, _in, _out): x = _in[0] #self.in_mask就是用来记录该channels层是否被prune了，为0则被prune了 self.in_mask += x.data.abs().sum(2, keepdim=True).sum(3, keepdim=True).cpu().sum(0, keepdim=True).view(-1) # 留下channels层，总和为0说明该channels被prune了 def _backward_hook(self, grad): #后向传播计算出gradient后执行的操作 self.out_mask += grad.data.abs().sum(2, keepdim=True).sum(3, keepdim=True).cpu().sum(0, keepdim=True).view(-1) # 留下channels层，总和为0说明该channels被prune了 new_grad = torch.ones_like(grad) return new_grad def forward(self, x): output = self.conv(x) noise = torch.zeros_like(output).normal_() output = output + noise #？ if self.training: output.register_hook(self._backward_hook) return output def melt(self): if self.conv.groups == 1: groups = 1 elif self.conv.groups == self.conv.out_channels: groups = int((self.out_mask != 0).sum()) else: assert False replacer = nn.Conv2d( in_channels = int((self.in_mask != 0).sum()), out_channels = int((self.out_mask != 0).sum()), kernel_size = self.conv.kernel_size, stride = self.conv.stride, padding = self.conv.padding, dilation = self.conv.dilation, groups = groups, bias = (self.conv.bias is not None) ).to(self.conv.weight.device) with torch.no_grad(): if self.conv.groups == 1: replacer.weight.set_(self.conv.weight[self.out_mask != 0][:, self.in_mask != 0]) else: replacer.weight.set_(self.conv.weight[self.out_mask != 0]) if self.conv.bias is not None: replacer.bias.set_(self.conv.bias[self.out_mask != 0]) return replacer @classmethod def transform(cls, net): r = [] def _inject(modules): keys = modules.keys() for k in keys: if len(modules[k]._modules) > 0: _inject(modules[k]._modules) if isinstance(modules[k], nn.Conv2d): modules[k] = Conv2dObserver(modules[k]) r.append(modules[k]) _inject(net._modules) return r
5）分类最后一层的全连接层怎么变：
class FinalLinearObserver(Meltable): ''' assert was in the last layer. only input was masked ''' def __init__(self, linear): super(FinalLinearObserver, self).__init__() assert isinstance(linear, nn.Linear) self.linear = linear self.in_mask = torch.zeros(linear.weight.shape[1]).to('cpu') self.f_hook = linear.register_forward_hook(self._forward_hook) #该linear层前向传播是进行的函数操作 def extra_repr(self): return '(%d, %d) -> (%d, %d)' % ( int(self.linear.weight.shape[1]), int(self.linear.weight.shape[0]), int((self.in_mask != 0).sum()), int(self.linear.weight.shape[0])) def _forward_hook(self, m, _in, _out): x = _in[0] self.in_mask += x.data.abs().cpu().sum(0, keepdim=True).view(-1) #列相加，每一列求和，输入的data中为0的列是因为那个channels被prune了， def forward(self, x): return self.linear(x) def melt(self): # 换成prune后的channels数 with torch.no_grad(): replacer = nn.Linear(int((self.in_mask != 0).sum()), self.linear.weight.shape[0]).to(self.linear.weight.device) replacer.weight.set_(self.linear.weight[:, self.in_mask != 0]) replacer.bias.set_(self.linear.bias) return replacer
这两个函数的作用在于将卷积层和全连接层分别封装成Conv2dObserver和FinalLinearObserver
如Conv2dObserver中就会有in_mask和out_mask两个参数，就是分别在训练的前向传播和后向传播中计算channels轴的和，最后为0则说明该轴已经被prune了，即：
def _forward_hook(self, m, _in, _out): x = _in[0] #self.in_mask就是用来记录该channels层是否被prune了，为0则被prune了 self.in_mask += x.data.abs().sum(2, keepdim=True).sum(3, keepdim=True).cpu().sum(0, keepdim=True).view(-1) # 留下channels层，总和为0说明该channels被prune了 def _backward_hook(self, grad): #后向传播计算出gradient后执行的操作 self.out_mask += grad.data.abs().sum(2, keepdim=True).sum(3, keepdim=True).cpu().sum(0, keepdim=True).view(-1) # 留下channels层，总和为0说明该channels被prune了 new_grad = torch.ones_like(grad) return new_grad
主要是用在下图标红部分：
即GBN变成剪枝后的BN的同时，卷积层和全连接层根据相邻的GBN层计算得到的in_mask和out_mask两个参数去剪枝对应的filter，令整个网络channels数是能链接起来的
6）gate的loss函数：
def get_gate_sparse_loss(masks, sparse_lambda): def _loss_hook(data, label, logits): loss = 0.0 for gbn in masks: if isinstance(gbn, GatedBatchNorm2d): loss += gbn.g.abs().sum() return sparse_lambda * loss return _loss_hook
这个是计算tock的损失的后半部分，后面看代码它是作为loss_hook的，的确是额外的loss
对应论文中的：
查看resnet56_prune中使用的resnet56的网络结构：
for name, module in pack.net.named_modules(): print(name) print(module)
返回：
DataParallel( (module): ResNet( (conv1): Conv2d(3, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False) (bn1): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) (layer1): Sequential( (0): BasicBlock( (conv1): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False) (bn1): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) (relu): ReLU(inplace=True) (conv2): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False) (bn2): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) (shortcut): Sequential() ) (1): BasicBlock( (conv1): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False) (bn1): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) (relu): ReLU(inplace=True) (conv2): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False) (bn2): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) (shortcut): Sequential() ) (2): BasicBlock( (conv1): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False) (bn1): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) (relu): ReLU(inplace=True) (conv2): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False) (bn2): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) (shortcut): Sequential() ) (3): BasicBlock( (conv1): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False) (bn1): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) (relu): ReLU(inplace=True) (conv2): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False) (bn2): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) (shortcut): Sequential() ) (4): BasicBlock( (conv1): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False) (bn1): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) (relu): ReLU(inplace=True) (conv2): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False) (bn2): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) (shortcut): Sequential() ) (5): BasicBlock( (conv1): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False) (bn1): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) (relu): ReLU(inplace=True) (conv2): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False) (bn2): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) (shortcut): Sequential() ) (6): BasicBlock( (conv1): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False) (bn1): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) (relu): ReLU(inplace=True) (conv2): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False) (bn2): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) (shortcut): Sequential() ) (7): BasicBlock( (conv1): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False) (bn1): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) (relu): ReLU(inplace=True) (conv2): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False) (bn2): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) (shortcut): Sequential() ) (8): BasicBlock( (conv1): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False) (bn1): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) (relu): ReLU(inplace=True) (conv2): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False) (bn2): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) (shortcut): Sequential() ) ) (layer2): Sequential( (0): BasicBlock( (conv1): Conv2d(16, 32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False) (bn1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) (relu): ReLU(inplace=True) (conv2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False) (bn2): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) (shortcut): Sequential( (0): Conv2d(16, 32, kernel_size=(1, 1), stride=(2, 2), bias=False) (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) ) ) ...
7）
剩下的Tick-Tock部分可能结合例子来讲比较好讲Gate Decorator: Global Filter Pruning Method for Accelerating Deep Convolutional Neural Networks - 模型压缩 - 3 - 代码学习，VGG16,Resnet