| 113 | |
| 114 | |
| 115 | class LayerScale_Block_CA(nn.Module): |
| 116 | # taken from https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py |
| 117 | # with slight modifications to add CA and LayerScale |
| 118 | def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0., |
| 119 | drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, Attention_block = Class_Attention, |
| 120 | Mlp_block=Mlp): |
| 121 | super().__init__() |
| 122 | self.norm1 = norm_layer(dim) |
| 123 | self.attn = Attention_block( |
| 124 | dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop) |
| 125 | self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() |
| 126 | self.norm2 = norm_layer(dim) |
| 127 | mlp_hidden_dim = int(dim * mlp_ratio) |
| 128 | self.mlp = Mlp_block(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) |
| 129 | |
| 130 | def forward(self, x, x_cls, attention=False, mask=None): |
| 131 | u = torch.cat((x_cls,x),dim=1) |
| 132 | if attention: |
| 133 | u_, cls_attn = self.attn(self.norm1(u), attention=True) |
| 134 | return cls_attn |
| 135 | else: |
| 136 | u_ = self.attn(self.norm1(u), mask=mask) |
| 137 | x_cls = x_cls + self.drop_path(u_) |
| 138 | x_cls = x_cls + self.drop_path(self.mlp(self.norm2(x_cls))) |
| 139 | return x_cls |
| 140 | |
| 141 | class Patch_Attention(nn.Module): |
| 142 | # taken from https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py |