from copy import deepcopy from typing import Optional, Union import torch from torch import nn from einops import rearrange, repeat from einops.layers.torch import Rearrange import tqdm from utils.dl.common.model import LayerActivation, get_model_device, get_model_size, set_module from .base import Abs, KTakesAll, ElasticDNNUtil, Layer_WrappedWithFBS from utils.common.log import logger class SqueezeLast(nn.Module): def __init__(self): super(SqueezeLast, self).__init__() def forward(self, x): return x.squeeze(-1) class ProjConv_WrappedWithFBS(Layer_WrappedWithFBS): def __init__(self, raw_conv2d: nn.Conv2d, r): super(ProjConv_WrappedWithFBS, self).__init__() self.fbs = nn.Sequential( Abs(), nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(raw_conv2d.in_channels, raw_conv2d.out_channels // r), nn.ReLU(), nn.Linear(raw_conv2d.out_channels // r, raw_conv2d.out_channels), nn.ReLU() ) self.raw_conv2d = raw_conv2d # self.raw_bn = raw_bn # remember clear the original BNs in the network nn.init.constant_(self.fbs[5].bias, 1.) nn.init.kaiming_normal_(self.fbs[5].weight) def forward(self, x): raw_x = self.raw_conv2d(x) if self.use_cached_channel_attention and self.cached_channel_attention is not None: channel_attention = self.cached_channel_attention else: self.cached_raw_channel_attention = self.fbs(x) self.cached_channel_attention = self.k_takes_all(self.cached_raw_channel_attention) channel_attention = self.cached_channel_attention return raw_x * channel_attention.unsqueeze(2).unsqueeze(3) class Linear_WrappedWithFBS(Layer_WrappedWithFBS): def __init__(self, linear: nn.Linear, r): super(Linear_WrappedWithFBS, self).__init__() self.linear = linear # for conv: (B, C_in, H, W) -> (B, C_in) -> (B, C_out) # for mlp in ViT: (B, #patches, D: dim of patches embedding) -> (B, D) -> (B, C_out) self.fbs = nn.Sequential( Rearrange('b n d -> b d n'), Abs(), nn.AdaptiveAvgPool1d(1), SqueezeLast(), nn.Linear(linear.in_features, linear.out_features // r), nn.ReLU(), nn.Linear(linear.out_features // r, linear.out_features), nn.ReLU() ) nn.init.constant_(self.fbs[6].bias, 1.) nn.init.kaiming_normal_(self.fbs[6].weight) def forward(self, x): if self.use_cached_channel_attention and self.cached_channel_attention is not None: channel_attention = self.cached_channel_attention else: self.cached_raw_channel_attention = self.fbs(x) self.cached_channel_attention = self.k_takes_all(self.cached_raw_channel_attention) channel_attention = self.cached_channel_attention raw_res = self.linear(x) return channel_attention.unsqueeze(1) * raw_res # class ToQKV_WrappedWithFBS(Layer_WrappedWithFBS): # """ # This regards to_q/to_k/to_v as a whole (in fact it consists of multiple heads) and prunes it. # It seems different channels of different heads are pruned according to the input. # This is different from "removing some head" or "removing the same channels in each head". # """ # def __init__(self, to_qkv: nn.Linear, r): # super(ToQKV_WrappedWithFBS, self).__init__() # # self.to_qkv = to_qkv # self.to_qk = nn.Linear(to_qkv.in_features, to_qkv.out_features // 3 * 2, bias=to_qkv.bias is not None) # self.to_v = nn.Linear(to_qkv.in_features, to_qkv.out_features // 3, bias=to_qkv.bias is not None) # self.to_qk.weight.data.copy_(to_qkv.weight.data[0: to_qkv.out_features // 3 * 2]) # if to_qkv.bias is not None: # self.to_qk.bias.data.copy_(to_qkv.bias.data[0: to_qkv.out_features // 3 * 2]) # self.to_v.weight.data.copy_(to_qkv.weight.data[to_qkv.out_features // 3 * 2: ]) # if to_qkv.bias is not None: # self.to_v.bias.data.copy_(to_qkv.bias.data[to_qkv.out_features // 3 * 2: ]) # self.fbs = nn.Sequential( # Rearrange('b n d -> b d n'), # Abs(), # nn.AdaptiveAvgPool1d(1), # SqueezeLast(), # nn.Linear(to_qkv.in_features, to_qkv.out_features // 3 // r), # nn.ReLU(), # # nn.Linear(to_qkv.out_features // 3 // r, to_qkv.out_features // 3), # nn.Linear(to_qkv.out_features // 3 // r, self.to_v.out_features), # nn.ReLU() # ) # nn.init.constant_(self.fbs[6].bias, 1.) # nn.init.kaiming_normal_(self.fbs[6].weight) # def forward(self, x): # if self.use_cached_channel_attention and self.cached_channel_attention is not None: # channel_attention = self.cached_channel_attention # else: # self.cached_raw_channel_attention = self.fbs(x) # # print() # # for attn in self.cached_raw_channel_attention.chunk(3, dim=1)[0: 1]: # # print(self.cached_raw_channel_attention.size(), attn.size()) # # print(self.k_takes_all.k) # # print(attn[0].nonzero(as_tuple=True)[0].size(), attn[0]) # self.cached_channel_attention = self.k_takes_all(self.cached_raw_channel_attention) # # for attn in self.cached_channel_attention.chunk(3, dim=1)[0: 1]: # # print(self.cached_channel_attention.size(), attn.size()) # # print(self.k_takes_all.k) # # print(attn[0].nonzero(as_tuple=True)[0].size(), attn[0]) # # print() # channel_attention = self.cached_channel_attention # qk = self.to_qk(x) # v = channel_attention.unsqueeze(1) * self.to_v(x) # return torch.cat([qk, v], dim=-1) # qkv = raw_res.chunk(3, dim = -1) # raw_v = qkv[2] # print('raw_k, raw_v', qkv[0].sum((0, 1))[0: 10], qkv[0].sum((0, 1)).nonzero(as_tuple=True)[0].size(), # qkv[1].sum((0, 1))[0: 10], qkv[1].sum((0, 1)).nonzero(as_tuple=True)[0].size(),) # print('raw_v', raw_v.size(), raw_v.sum((0, 1))[0: 10], raw_v.sum((0, 1)).nonzero(as_tuple=True)[0].size()) # qkv_attn = channel_attention.chunk(3, dim=-1) # print('attn', [attn[0][0: 10] for attn in qkv_attn]) # print(channel_attention.unsqueeze(1).size(), raw_res.size()) # print('fbs', channel_attention.size(), raw_res.size()) # return channel_attention.unsqueeze(1) * raw_res class LinearStaticFBS(nn.Module): def __init__(self, static_channel_attention): super(LinearStaticFBS, self).__init__() assert static_channel_attention.dim() == 2 and static_channel_attention.size(0) == 1 self.static_channel_attention = nn.Parameter(static_channel_attention, requires_grad=False) # (1, dim) def forward(self, x): # print('staticfbs', x, self.static_channel_attention.unsqueeze(1)) return x * self.static_channel_attention.unsqueeze(1) from .cnn import StaticFBS as ConvStaticFBS class ElasticViTUtil(ElasticDNNUtil): def convert_raw_dnn_to_master_dnn(self, raw_dnn: nn.Module, r: float, ignore_layers=[]): assert len(ignore_layers) == 0, 'not supported yet' raw_vit = deepcopy(raw_dnn) set_module(raw_vit, 'patch_embed.proj', ProjConv_WrappedWithFBS(raw_vit.patch_embed.proj, r)) for name, module in raw_vit.named_modules(): if name.endswith('mlp'): set_module(module, 'fc1', Linear_WrappedWithFBS(module.fc1, r)) return raw_vit # def set_master_dnn_sparsity(self, master_dnn: nn.Module, sparsity: float): # for name, module in master_dnn.named_modules(): # if not name.endswith('attn'): # continue # q_features = module.qkv.to_qk.out_features // 2 # if (q_features - int(q_features * sparsity)) % module.num_heads != 0: # # tune sparsity to ensure #unpruned channel % num_heads == 0 # # so that the pruning seems to reduce the dim_head of each head # tuned_sparsity = 1. - int((q_features - int(q_features * sparsity)) / module.num_heads) * module.num_heads / q_features # logger.debug(f'tune sparsity from {sparsity:.2f} to {tuned_sparsity}') # sparsity = tuned_sparsity # break # return super().set_master_dnn_sparsity(master_dnn, sparsity) def select_most_rep_sample(self, master_dnn: nn.Module, samples: torch.Tensor): return samples[0].unsqueeze(0) def extract_surrogate_dnn_via_samples(self, master_dnn: nn.Module, samples: torch.Tensor): sample = self.select_most_rep_sample(master_dnn, samples) assert sample.dim() == 4 and sample.size(0) == 1 print('WARN: for debug, modify cls_token and pos_embed') master_dnn.pos_embed.data = torch.zeros_like(master_dnn.pos_embed.data) print('before') master_dnn.eval() self.clear_cached_channel_attention_in_master_dnn(master_dnn) # debug: add hooks hooks = { 'blocks_input': LayerActivation(master_dnn.blocks, True, 'cuda') } with torch.no_grad(): master_dnn_output = master_dnn(sample) for k, v in hooks.items(): print(f'{k}: {v.input.size()}') print('after') boosted_vit = master_dnn def get_unpruned_indexes_from_channel_attn(channel_attn: torch.Tensor, k): assert channel_attn.size(0) == 1, 'use A representative sample to generate channel attentions' res = channel_attn[0].nonzero(as_tuple=True)[0] return res proj = boosted_vit.patch_embed.proj proj_unpruned_indexes = get_unpruned_indexes_from_channel_attn( proj.cached_channel_attention, proj.k_takes_all.k) # 1.1 prune proj itself proj_conv = proj.raw_conv2d new_proj = nn.Conv2d(proj_conv.in_channels, proj_unpruned_indexes.size(0), proj_conv.kernel_size, proj_conv.stride, proj_conv.padding, proj_conv.dilation, proj_conv.groups, proj_conv.bias is not None, proj_conv.padding_mode, proj_conv.weight.device) new_proj.weight.data.copy_(proj_conv.weight.data[proj_unpruned_indexes]) if new_proj.bias is not None: new_proj.bias.data.copy_(proj_conv.bias.data[proj_unpruned_indexes]) set_module(boosted_vit.patch_embed, 'proj', nn.Sequential(new_proj, ConvStaticFBS(proj.cached_channel_attention[0][proj_unpruned_indexes]))) # print(boosted_vit.pos_embed.size(), boosted_vit.cls_token.size()) boosted_vit.pos_embed.data = boosted_vit.pos_embed.data[:, :, proj_unpruned_indexes] boosted_vit.cls_token.data = boosted_vit.cls_token.data[:, :, proj_unpruned_indexes] def reduce_linear_output(raw_linear: nn.Linear, layer_name, unpruned_indexes: torch.Tensor): new_linear = nn.Linear(raw_linear.in_features, unpruned_indexes.size(0), raw_linear.bias is not None) new_linear.weight.data.copy_(raw_linear.weight.data[unpruned_indexes]) if raw_linear.bias is not None: new_linear.bias.data.copy_(raw_linear.bias.data[unpruned_indexes]) set_module(boosted_vit, layer_name, new_linear) def reduce_linear_input(raw_linear: nn.Linear, layer_name, unpruned_indexes: torch.Tensor): new_linear = nn.Linear(unpruned_indexes.size(0), raw_linear.out_features, raw_linear.bias is not None) new_linear.weight.data.copy_(raw_linear.weight.data[:, unpruned_indexes]) if raw_linear.bias is not None: new_linear.bias.data.copy_(raw_linear.bias.data) set_module(boosted_vit, layer_name, new_linear) def reduce_norm(raw_norm: nn.LayerNorm, layer_name, unpruned_indexes: torch.Tensor): new_norm = nn.LayerNorm(unpruned_indexes.size(0), raw_norm.eps, raw_norm.elementwise_affine) new_norm.weight.data.copy_(raw_norm.weight.data[unpruned_indexes]) new_norm.bias.data.copy_(raw_norm.bias.data[unpruned_indexes]) set_module(boosted_vit, layer_name, new_norm) # 1.2 prune blocks.x.norm1/to_qkv/proj/fc1/fc2 for block_i, block in enumerate(boosted_vit.blocks): attn = block.attn ff = block.mlp reduce_norm(block.norm1, f'blocks.{block_i}.norm1', proj_unpruned_indexes) reduce_linear_input(attn.qkv, f'blocks.{block_i}.attn.qkv', proj_unpruned_indexes) reduce_linear_output(attn.proj, f'blocks.{block_i}.attn.proj', proj_unpruned_indexes) reduce_norm(block.norm2, f'blocks.{block_i}.norm2', proj_unpruned_indexes) reduce_linear_input(ff.fc1.linear, f'blocks.{block_i}.mlp.fc1.linear', proj_unpruned_indexes) reduce_linear_output(ff.fc2, f'blocks.{block_i}.mlp.fc2', proj_unpruned_indexes) # 1.3 prune norm, head reduce_norm(boosted_vit.norm, f'norm', proj_unpruned_indexes) reduce_linear_input(boosted_vit.head, f'head', proj_unpruned_indexes) # 2. prune blocks.x.fc1 for block_i, block in enumerate(boosted_vit.blocks): attn = block.attn ff = block.mlp fc1 = ff.fc1 fc1_unpruned_indexes = get_unpruned_indexes_from_channel_attn(fc1.cached_channel_attention, fc1.k_takes_all.k) fc1_linear = fc1.linear new_linear = nn.Linear(fc1_linear.in_features, fc1_unpruned_indexes.size(0), fc1_linear.bias is not None) new_linear.weight.data.copy_(fc1_linear.weight.data[fc1_unpruned_indexes]) if fc1_linear.bias is not None: new_linear.bias.data.copy_(fc1_linear.bias.data[fc1_unpruned_indexes]) set_module(boosted_vit, f'blocks.{block_i}.mlp.fc1', nn.Sequential(new_linear, LinearStaticFBS(fc1.cached_channel_attention[:, fc1_unpruned_indexes]))) reduce_linear_input(ff.fc2, f'blocks.{block_i}.mlp.fc2', fc1_unpruned_indexes) surrogate_dnn = boosted_vit surrogate_dnn.eval() surrogate_dnn = surrogate_dnn.to(get_model_device(master_dnn)) print(surrogate_dnn) hooks = { 'blocks_input': LayerActivation(surrogate_dnn.blocks, True, 'cuda') } with torch.no_grad(): surrogate_dnn_output = surrogate_dnn(sample) for k, v in hooks.items(): print(f'{k}: {v.input.size()}') output_diff = ((surrogate_dnn_output - master_dnn_output) ** 2).sum() assert output_diff < 1e-4, output_diff logger.info(f'output diff of master and surrogate DNN: {output_diff}') return boosted_vit