| 374 | return self.pe[:seq_len, :].unsqueeze(0) |
| 375 | |
| 376 | class Transformer(nn.Module): |
| 377 | def __init__(self, width, layers, heads, window_size=None, block_cls=ResAttBlock, drop_path_rate=0.0): |
| 378 | super().__init__() |
| 379 | self.width = width |
| 380 | self.layers = layers |
| 381 | blocks = [] |
| 382 | dpr = [x.item() for x in torch.linspace(0, drop_path_rate, layers)] # stochastic depth decay rule |
| 383 | inter_dpr = [0.0] + dpr |
| 384 | if drop_path_rate > 0.0: |
| 385 | print(f"inter_dpr: {inter_dpr}") |
| 386 | for _ in range(layers): |
| 387 | layer = block_cls(width, heads, window_size=window_size, drop_path_rate=inter_dpr[_]) |
| 388 | blocks.append(layer) |
| 389 | |
| 390 | self.resblocks = nn.Sequential(*blocks) |
| 391 | self.grad_checkpointing = False |
| 392 | |
| 393 | def set_grad_checkpointing(self, flag=True): |
| 394 | self.grad_checkpointing = flag |
| 395 | |
| 396 | def forward(self, x, condition=None): |
| 397 | for res_i, module in enumerate(self.resblocks): |
| 398 | if self.grad_checkpointing: |
| 399 | x = checkpoint(module, x, res_i, condition, use_reentrant=False) |
| 400 | else: |
| 401 | x = module(x, res_i, condition) |
| 402 | |
| 403 | return x |
| 404 | |
| 405 | class TransformerBase(nn.Module): |
| 406 | def __init__(self, width, layers, heads, window_size, token_len, block_cls, drop_path_rate=0.0): |