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