| 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 |