<gh_stars>0
import torch
import torch.nn as nn
import numpy as np

class Attention(nn.Module):
    def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.):
        super().__init__()
        self.num_heads = num_heads
        head_dim = dim // num_heads
        self.scale = qk_scale or head_dim ** -0.5

        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)

    def forward(self, x):
        B, N, C = x.shape
        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
        q, k, v = qkv[0], qkv[1], qkv[2]   # make torchscript happy (cannot use tensor as tuple)

        attn = (q @ k.transpose(-2, -1)) * self.scale
        attn = attn.softmax(dim=-1)
        attn = self.attn_drop(attn)

        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
        x = self.proj(x)
        x = self.proj_drop(x)
        return x

class PrunableAttention(Attention):
    def __init__(self, attn_module):
        super().__init__(attn_module.qkv.in_features, attn_module.num_heads, True, attn_module.scale, attn_module.attn_drop.p, attn_module.proj_drop.p)
        self.is_pruned = False
        self.num_gates = attn_module.qkv.in_features // self.num_heads
        self.zeta = nn.Parameter(torch.rand(1, 1, self.num_heads, 1, self.num_gates) * 0.01)
        self.pruned_zeta = torch.ones_like(self.zeta)
        beta=1.
        gamma=2.
        for n, x in zip(('beta', 'gamma'), (torch.tensor([x], requires_grad=False) for x in (beta, gamma))):
            self.register_buffer(n, x)  # self.beta will be created (same for gamma)         
    
    def forward(self, x):
        B, N, C = x.shape
        z = self.pruned_zeta if self.is_pruned else self.get_zeta_t()
        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4) # 3, B, H, N, d(C/H)
        qkv *= z
        q, k, v = qkv[0], qkv[1], qkv[2]   # make torchscript happy (cannot use tensor as tuple) # B, H, N, d

        attn = (q @ k.transpose(-2, -1)) * self.scale
        attn = attn.softmax(dim=-1)
        attn = self.attn_drop(attn)

        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
        x *= z.reshape(-1)[None, None, :]
        x = self.proj(x)
        x = self.proj_drop(x)
        return x
    
    def get_zeta_i(self):
        return self.__generalized_logistic(self.zeta)
    
    def get_zeta_t(self):
        zeta_i = self.get_zeta_i()
        return self.__continous_heavy_side(zeta_i)

    def set_beta_gamma(self, beta, gamma):
        self.beta.data.copy_(torch.Tensor([beta]))
        self.gamma.data.copy_(torch.Tensor([gamma]))
      
    def __generalized_logistic(self, x):
        return 1./(1.+torch.exp(-self.beta*x))
    
    def __continous_heavy_side(self, x):
        return 1-torch.exp(-self.gamma*x)+x*torch.exp(-self.gamma)
    
    def prune(self, threshold):
        self.is_pruned = True
        self.pruned_zeta = (self.get_zeta_t()>threshold).float()
        # if self.is_imp and self.pruned_zeta.sum()==0:
        #     self.pruned_zeta[torch.argmax(self.get_zeta_t()).item()] = 1.
        self.zeta.requires_grad = False

    def unprune(self):
        self.is_pruned = False
        self.zeta.requires_grad = True

    def get_params_count(self):
        dim = self.qkv.in_features
        active = self.pruned_zeta.sum().data
        total_params = dim*dim*3 + dim*3
        total_params += dim*dim + dim
        active_params = dim*active*3 + active*3
        active_params += active*dim +dim
        return total_params, active_params
    
    def get_flops(self, num_patches):
        H = self.num_heads
        N = num_patches
        d = self.num_gates
        sd = self.pruned_zeta.sum()
        ssd = np.array([(self.pruned_zeta[0,0,i].sum().item())**2 for i in range(self.num_heads)]).sum()
        total_flops = N * (H*d * (3*H*d)) + 3*N*H*d #linear: qkv
        total_flops += H*N*d*N + H*N*N #q@k
        total_flops += 5*H*N*N #softmax
        total_flops += H*N*N*d #attn@v
        total_flops += N * (H*d * (H*d)) + N*H*d #linear: proj
        
        active_flops = N * (H*d * (3*sd)) + 3*N*sd #linear: qkv
        active_flops += N*N*sd + H*N*N #q@k
        active_flops += 5*H*N*N #softmax
        active_flops += N*N*sd #attn@v
        active_flops += N * (sd * (H*d)) + N*H*d #linear: proj
        return total_flops, active_flops

    @staticmethod
    def from_attn(attn_module):
        attn_module = PrunableAttention(attn_module)
        return attn_module


class ModuleInjection:
    pruning_method = 'full'
    prunable_modules = []

    @staticmethod
    def make_prunable(attn_module):
        """Make a attn_module prunable.
        :param attn_module: A Attention module
        :return: a attn_module that can be trained to
        """
        if ModuleInjection.pruning_method == 'full':
            return attn_module
        attn_module = PrunableAttention.from_attn(attn_module)
        ModuleInjection.prunable_modules.append(attn_module)
        return attn_module