| 150 | |
| 151 | |
| 152 | class DropPath(nn.Module): |
| 153 | def __init__(self, drop_prob=None): |
| 154 | super().__init__() |
| 155 | self.drop_prob = float(drop_prob) |
| 156 | |
| 157 | def forward(self, x): |
| 158 | batch, drop_prob, device, dtype = x.shape[0], self.drop_prob, x.device, x.dtype |
| 159 | |
| 160 | if drop_prob <= 0. or not self.training: |
| 161 | return x |
| 162 | |
| 163 | keep_prob = 1 - self.drop_prob |
| 164 | shape = (batch, *((1,) * (x.ndim - 1))) |
| 165 | |
| 166 | keep_mask = torch.zeros(shape, device=device).float().uniform_(0, 1) < keep_prob |
| 167 | output = x.div(keep_prob) * keep_mask.float() |
| 168 | return output |
| 169 | |
| 170 | |
| 171 | class Tokenizer(nn.Module): |