import torch from terminaltables import AsciiTable from copy import deepcopy import numpy as np import torch.nn.functional as F def get_sr_flag(epoch, sr): # return epoch >= 5 and sr return sr def parse_module_defs(module_defs): CBL_idx = [] Conv_idx = [] for i, module_def in enumerate(module_defs): # 添加了 or module_def['type'] == 'ds_conv' 将ds_conv也纳入剪枝范围，如果不考虑的话就在下面的不需要剪枝那里添加 if module_def['type'] == 'convolutional' or module_def['type'] == 'ds_conv': if module_def['batch_normalize'] == '1': CBL_idx.append(i) #CBL_idx= 即为有卷积也有bn的层的idx #[0, 1, 2, 3, 5, 6, 7, 9, 10, 12, 13, 14, 16, 17, 19, 20, 22, 23, 25, 26, 28, 29, 31, 32, 34, 35, 37, 38, # 39, 41, 42, 44, 45, 47, 48, 50, 51, 53, 54, 56, 57, 59, 60, 62, 63, 64, 66, 67, 69, 70, 72, 73, 75, 76, 77, # 78, 79, 80, 84, 87, 88, 89, 90, 91, 92, 96, 99, 100, 101, 102, 103, 104] else: #Conv_idx = [81, 93, 105] 即为有卷积没有bn的层的idx Conv_idx.append(i) ignore_idx = set() #哪些层不需要剪枝 for i, module_def in enumerate(module_defs): #将ds_conv纳入不剪枝范围 if module_def['type'] == 'ds_conv': ignore_idx.add(i-1) if module_def['type'] == 'shortcut': ignore_idx.add(i-1) identity_idx = (i + int(module_def['from'])) if module_defs[identity_idx]['type'] == 'convolutional': ignore_idx.add(identity_idx) elif module_defs[identity_idx]['type'] == 'shortcut': ignore_idx.add(identity_idx - 1) ignore_idx.add(84) ignore_idx.add(96) # print(ignore_idx) #ignore_idx={1, 3, 5, 7, 10, 12, 14, 17, 20, 23, 26, 29, 32, 35, 37, 39, 42, 45, 48, 51, 54, 57, 60, 62, 64, 67, 70, 73, 84, 96} 改结构前 # ignore_idx={1, 3, 5, 7, 10, 12, 13, 14, 16, 17, 19, 20, 22, 23, 25, 26, 28, 29, 31, 32, 34, 35, 37, 38, 39, 41, 42, 44, 45, 47, # 48, 50, 51, 53, 54, 56, 57, 59, 60, 62, 63, 64, 66, 67, 69, 70, 72, 73, 84, 96} prune_idx = [idx for idx in CBL_idx if idx not in ignore_idx] #prune_idx =[0, 2, 6, 9, 13, 16, 19, 22, 25, 28, 31, 34, 38, 41, 44, 47, 50, 53, 56, 59, 63, 66, 69, 72, 75, # 76, 77, 78, 79, 80, 87, 88, 89, 90, 91, 92, 99, 100, 101, 102, 103, 104] 改结构前 # prune_idx = [0, 2, 6, 9, 75, 76, 77, 78, 79, 80, 87, 88, 89, 90, 91, 92, 99, 100, 101, 102, 103, 104] # 返回CBL组件的id，单独的Conv层的id，以及需要被剪枝的层的id return CBL_idx, Conv_idx, prune_idx def gather_bn_weights(module_list, prune_idx): size_list = [module_list[idx][1].weight.data.shape[0] for idx in prune_idx] # 存储prune_idx对应层的 filter数量 # print('module_list[0][1].weight.data.shape[0] = ',module_list[0][1].weight.data.shape[0]) # print("size_list:",size_list) # ds_conv对应的size_list [32, 32, 64, 64, 512, 1024, 512, 1024, 512, 1024, 256, 512, 256, 512, 256, 512, 128, 256, 128, 256, 128, 256] #[32, 32, 64, 64, 128, 128, 128, 128, 128, 128, 128, 128, 256, 256, 256, 256, 256, 256, 256, 256, 512, 512, 512, 512, 512, # 1024, 512, 1024, 512, 1024, 256, 512, 256, 512, 256, 512, 128, 256, 128, 256, 128, 256] bn_weights = torch.zeros(sum(size_list)) # print('bn_weights = ', bn_weights) # bn_weights == tensor([0., 0., 0., ..., 0., 0., 0.]) index = 0 for idx, size in zip(prune_idx, size_list): # print('idx = ' ,idx, ' | size = ',size) bn_weights[index:(index + size)] = module_list[idx][1].weight.data.abs().clone() index += size # print('bn_weights = ' ,bn_weights, ' | index = ',index) # print('module_list[0][1].weight.data.abs().clone()_len = ',len(module_list[0][1].weight.data.abs().clone())) # 获取CBL组件的BN层的权重，即Gamma参数，我们会根据这个参数来剪枝 return bn_weights def write_cfg(cfg_file, module_defs): with open(cfg_file, 'w') as f: for module_def in module_defs: f.write(f"[{module_def['type']}]\n") for key, value in module_def.items(): if key != 'type': f.write(f"{key}={value}\n") f.write("\n") return cfg_file class BNOptimizer(): @staticmethod def updateBN(sr_flag, module_list, s, prune_idx): if sr_flag: for idx in prune_idx: # Squential(Conv, BN, Lrelu) bn_module = module_list[idx][1] # 将批标准化模块(bn_module)的权重矩阵的梯度乘以一个缩放因子后添加到该权重矩阵的梯度上。 bn_module.weight.grad.data.add_(s * torch.sign(bn_module.weight.data)) # L1 def obtain_quantiles(bn_weights, num_quantile=5): sorted_bn_weights, i = torch.sort(bn_weights) total = sorted_bn_weights.shape[0] quantiles = sorted_bn_weights.tolist()[-1::-total//num_quantile][::-1] print("\nBN weights quantile:") quantile_table = [ [f'{i}/{num_quantile}' for i in range(1, num_quantile+1)], ["%.3f" % quantile for quantile in quantiles] ] print(AsciiTable(quantile_table).table) return quantiles def get_input_mask(module_defs, idx, CBLidx2mask): # print("CBLidx2mask ： ", CBLidx2mask) if idx == 0: return np.ones(3) if module_defs[idx - 1]['type'] == 'convolutional': return CBLidx2mask[idx - 1] elif module_defs[idx - 1]['type'] == 'ds_conv': return CBLidx2mask[idx - 1] elif module_defs[idx - 1]['type'] == 'shortcut': return CBLidx2mask[idx - 2] elif module_defs[idx - 1]['type'] == 'route': route_in_idxs = [] for layer_i in module_defs[idx - 1]['layers'].split(","): if int(layer_i) < 0: route_in_idxs.append(idx - 1 + int(layer_i)) else: route_in_idxs.append(int(layer_i)) if len(route_in_idxs) == 1: return CBLidx2mask[route_in_idxs[0]] elif len(route_in_idxs) == 2: return np.concatenate([CBLidx2mask[in_idx - 1] for in_idx in route_in_idxs]) else: print("Something wrong with route module!") raise Exception def init_weights_from_loose_model(compact_model, loose_model, CBL_idx, Conv_idx, CBLidx2mask): # (14): Sequential( # (ds_conv_d_14): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=128) # (ds_conv_p_14): Conv2d(128, 256, kernel_size=(1, 1), stride=(1, 1)) # (batch_norm_14): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) # (leaky_14): LeakyReLU(negative_slope=0.1, inplace=True) # ) for idx in CBL_idx: compact_CBL = compact_model.module_list[idx] loose_CBL = loose_model.module_list[idx] # print("idx:",idx,"--compact_CBL :", compact_CBL, "--compact_CBL[1] :", compact_CBL[1],"--compact_CBL[2] :", compact_CBL[2]) # print("idx:",idx,"--loose_CBL :", loose_CBL, "--loose_CBL[1] :", loose_CBL[1],"--loose_CBL[2] :", loose_CBL[2]) # np.argwhere返回非0元素的索引，X[:,0]是numpy中数组的一种写法，表示对一个二维数组，取该二维数组第一维中的所有数据，第二维中取第0个数据 out_channel_idx = np.argwhere(CBLidx2mask[idx])[:, 0].tolist() # print("idx:",idx,"--out_channel_idx:", len(out_channel_idx)) # if (idx == 14): # print("idx:",idx,"--out_channel_idx:", out_channel_idx, "--len:",len(out_channel_idx)) # print("CBLidx2mask[idx] :",CBLidx2mask[idx],"--len:",len(CBLidx2mask[idx])) # print("loose_model.module_list[idx][1]:",loose_model.module_list[idx][1]) # 获取剪枝后的模型当前BN层的权重 if(compact_model.module_defs[idx]['type'] == 'ds_conv') : compact_bn, loose_bn = compact_CBL[2], loose_CBL[2] else: compact_bn, loose_bn = compact_CBL[1], loose_CBL[1] # if idx==14: # print("compact_bn.weight.data:",compact_bn.weight.data,"--len",len(compact_bn.weight.data)) # print("loose_bn.weight.data:",loose_bn.weight.data,"--len",len(loose_bn.weight.data)) # print("loose_bn.weight.data[out_channel_idx]:", loose_bn.weight.data[out_channel_idx],"--len",len(loose_bn.weight.data[out_channel_idx])) compact_bn.weight.data = loose_bn.weight.data[out_channel_idx].clone() compact_bn.bias.data = loose_bn.bias.data[out_channel_idx].clone() compact_bn.running_mean.data = loose_bn.running_mean.data[out_channel_idx].clone() compact_bn.running_var.data = loose_bn.running_var.data[out_channel_idx].clone() # 获取剪枝后的模型当前卷积层的权重，这和上一个卷积层的剪枝情况有关 # print("idx:",idx) input_mask = get_input_mask(loose_model.module_defs, idx, CBLidx2mask) in_channel_idx = np.argwhere(input_mask)[:, 0].tolist() # if (idx == 14): # print("idx:",idx,"--in_channel_idx:", in_channel_idx, "--len:",len(in_channel_idx)) # print("idx:",idx,"--in_channel_idx:",len(in_channel_idx)) if(compact_model.module_defs[idx]['type'] == 'ds_conv') : compact_conv, loose_conv = compact_CBL[0], loose_CBL[0] # compact_conv, loose_conv = compact_CBL[1], loose_CBL[1] compact_conv2, loose_conv2 = compact_CBL[1], loose_CBL[1] # 拷贝权重到剪枝后的模型中去 # print("loose_conv2.weight.data:",loose_conv2.weight.data,"--len:",len(loose_conv2.weight.data)) # print("loose_conv2.weight.data[:, in_channel_idx, :, :]:",loose_conv2.weight.data[:, in_channel_idx, :, :],"--len:",len(loose_conv2.weight.data[:, in_channel_idx, :, :])) # print("compact_conv.weight.data:",compact_conv.weight.data,"--len:",len(compact_conv.weight.data)) # print("compact_conv2.weight.data:",compact_conv2.weight.data,"--len:",len(compact_conv2.weight.data)) # print("loose_conv.weight.data:",loose_conv.weight.data,"--len:",len(loose_conv.weight.data)) # print("loose_conv.weight.data[:, in_channel_idx, :, :]:",loose_conv.weight.data[:, in_channel_idx, :, :],"--len:",len(loose_conv.weight.data[:, in_channel_idx, :, :])) # tmp1 = loose_conv.weight.data[:, in_channel_idx, :, :].clone() # compact_conv.weight.data = tmp1[in_channel_idx, :, :, :].clone() compact_conv.weight.data = loose_conv.weight.data.clone() compact_conv.bias.data = loose_conv.bias.data.clone() # tmp2 = loose_conv2.weight.data[:, in_channel_idx, :, :].clone() # compact_conv2.weight.data = tmp2[out_channel_idx, :, :, :].clone() compact_conv2.weight.data = loose_conv2.weight.data.clone() compact_conv2.bias.data = loose_conv2.bias.data.clone() else: compact_conv, loose_conv = compact_CBL[0], loose_CBL[0] # print("idx:",idx,"--compact_CBL[0]:",compact_CBL[0],"--loose_CBL[0]",loose_CBL[0]) #还有一个可能就是ds_conv中有两个卷积层都需要复制权重！ # print("idx:",idx,"--compact_CBL[1]:",compact_CBL[1],"--loose_CBL[0]",loose_CBL[1]) # 拷贝权重到剪枝后的模型中去 tmp = loose_conv.weight.data[:, in_channel_idx, :, :].clone() compact_conv.weight.data = tmp[out_channel_idx, :, :, :].clone() for idx in Conv_idx: compact_conv = compact_model.module_list[idx][0] loose_conv = loose_model.module_list[idx][0] # 虽然当前层是不带BN的卷积层，但仍然和上一个层的剪枝情况是相关的 input_mask = get_input_mask(loose_model.module_defs, idx, CBLidx2mask) in_channel_idx = np.argwhere(input_mask)[:, 0].tolist() # 拷贝权重到剪枝后的模型中去 compact_conv.weight.data = loose_conv.weight.data[:, in_channel_idx, :, :].clone() compact_conv.bias.data = loose_conv.bias.data.clone() def prune_model_keep_size(model, prune_idx, CBL_idx, CBLidx2mask): # 先拷贝一份原始的模型参数 pruned_model = deepcopy(model) # 对需要剪枝的层分别处理 for idx in prune_idx: # 需要保留的通道 mask = torch.from_numpy(CBLidx2mask[idx]).cuda() # 获取BN层的gamma参数，即BN层的权重 bn_module = pruned_model.module_list[idx][1] # print("bn_module:", pruned_model.module_list[idx][1]) bn_module.weight.data.mul_(mask) # 获取保留下来的通道产生的激活值，注意是每个通道分别获取的 activation = F.leaky_relu((1 - mask) * bn_module.bias.data, 0.1) # 两个上采样层前的卷积层 next_idx_list = [idx + 1] if idx == 79: next_idx_list.append(84) elif idx == 91: next_idx_list.append(96) # print("idx:",idx,"--next_idx_list:",next_idx_list) # 对下一层进行处理 for next_idx in next_idx_list: # 当前层的BN剪枝之后会对下一个卷积层造成影响 next_conv = pruned_model.module_list[next_idx][0] # dim=(2,3)即在(w,h)维度上进行求和，因为是通道剪枝，一个通道对应着(w,h)这个矩形 conv_sum = next_conv.weight.data.sum(dim=(2, 3)) # 将卷积层的权重和激活值相乘获得剪枝后的每个通道的偏置，以更新下一个BN层或者下一个带偏置的卷积层的偏执（因为单独的卷积层是不会被剪枝的，所以只对偏置有影响 # print("idx:", {idx} , "| next_idx:" , {next_idx}) # print("conv_sum :", conv_sum) # print("activation :", activation) # print("--------------------------") offset = conv_sum.matmul(activation.reshape(-1, 1)).reshape(-1) if next_idx in CBL_idx: next_bn = pruned_model.module_list[next_idx][1] next_bn.running_mean.data.sub_(offset) else: next_conv.bias.data.add_(offset) bn_module.bias.data.mul_(mask) # 返回剪枝后的模型 return pruned_model def obtain_bn_mask(bn_module, thre): thre = thre.cuda() # ge(a, b)相当于 a>= b mask = bn_module.weight.data.abs().ge(thre).float() # print('thre = ',thre,"| mask = ", mask) # 返回通道是否需要剪枝的通道状态 return mask