(self, x)
| 43 | self.exclude_first_token = exclude_first_token # exclude CLS token |
| 44 | |
| 45 | def forward(self, x): |
| 46 | if not self.training or self.prob == 0.: |
| 47 | return x |
| 48 | |
| 49 | if self.exclude_first_token: |
| 50 | cls_tokens, x = x[:, :1], x[:, 1:] |
| 51 | else: |
| 52 | cls_tokens = torch.jit.annotate(torch.Tensor, x[:, :1]) |
| 53 | |
| 54 | batch = x.size()[0] |
| 55 | num_tokens = x.size()[1] |
| 56 | |
| 57 | batch_indices = torch.arange(batch) |
| 58 | batch_indices = batch_indices[..., None] |
| 59 | |
| 60 | keep_prob = 1 - self.prob |
| 61 | num_patches_keep = max(1, int(num_tokens * keep_prob)) |
| 62 | |
| 63 | rand = torch.randn(batch, num_tokens) |
| 64 | patch_indices_keep = rand.topk(num_patches_keep, dim=-1).indices |
| 65 | |
| 66 | x = x[batch_indices, patch_indices_keep] |
| 67 | |
| 68 | if self.exclude_first_token: |
| 69 | x = torch.cat((cls_tokens, x), dim=1) |
| 70 | |
| 71 | return x |
nothing calls this directly
no outgoing calls
no test coverage detected