| 70 | |
| 71 | |
| 72 | class Class_Attention(nn.Module): |
| 73 | # taken from https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py |
| 74 | # with slight modifications to do CA |
| 75 | def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.): |
| 76 | super().__init__() |
| 77 | self.num_heads = num_heads |
| 78 | head_dim = dim // num_heads |
| 79 | self.scale = qk_scale or head_dim ** -0.5 |
| 80 | |
| 81 | self.q = nn.Linear(dim, dim, bias=qkv_bias) |
| 82 | self.k = nn.Linear(dim, dim, bias=qkv_bias) |
| 83 | self.v = nn.Linear(dim, dim, bias=qkv_bias) |
| 84 | self.attn_drop = nn.Dropout(attn_drop) |
| 85 | self.proj = nn.Linear(dim, dim) |
| 86 | self.proj_drop = nn.Dropout(proj_drop) |
| 87 | |
| 88 | |
| 89 | def forward(self, x, attention=False, mask=None): |
| 90 | |
| 91 | B, N, C = x.shape |
| 92 | q = self.q(x[:,0]).unsqueeze(1).reshape(B, 1, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3) |
| 93 | k = self.k(x).reshape(B, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3) |
| 94 | |
| 95 | q = q * self.scale |
| 96 | v = self.v(x).reshape(B, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3) |
| 97 | |
| 98 | attn = (q @ k.transpose(-2, -1)) |
| 99 | if mask is not None: |
| 100 | mask_temp = torch.cat([torch.ones(B,1).bool().cuda(), mask],dim=1).unsqueeze(1).unsqueeze(1).expand(-1,self.num_heads,-1,-1) |
| 101 | attn = attn.masked_fill_(~mask_temp.bool(), float("-inf")) |
| 102 | attn = attn.softmax(dim=-1) |
| 103 | attn = self.attn_drop(attn) |
| 104 | |
| 105 | x_cls = (attn @ v).transpose(1, 2).reshape(B, 1, C) |
| 106 | x_cls = self.proj(x_cls) |
| 107 | x_cls = self.proj_drop(x_cls) |
| 108 | |
| 109 | if attention: |
| 110 | return x_cls, attn |
| 111 | else: |
| 112 | return x_cls |
| 113 | |
| 114 | |
| 115 | class LayerScale_Block_CA(nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected