| 298 | |
| 299 | |
| 300 | class TransformerBlock(nn.Module): |
| 301 | def __init__( |
| 302 | self, |
| 303 | width: int, |
| 304 | layers: int, |
| 305 | heads: int, |
| 306 | mlp_ratio: float = 4.0, |
| 307 | act_layer: Callable = nn.GELU, |
| 308 | norm_layer: Callable = nn.LayerNorm, |
| 309 | ): |
| 310 | super().__init__() |
| 311 | self.width = width |
| 312 | self.layers = layers |
| 313 | |
| 314 | self.resblocks = nn.ModuleList([ |
| 315 | VisualAttentionBlock( |
| 316 | width, heads, mlp_ratio, act_layer=act_layer, norm_layer=norm_layer) |
| 317 | for _ in range(layers) |
| 318 | ]) |
| 319 | |
| 320 | def get_cast_dtype(self) -> torch.dtype: |
| 321 | return self.resblocks[0].mlp.c_fc.weight.dtype |
| 322 | |
| 323 | def get_cast_device(self) -> torch.device: |
| 324 | return self.resblocks[0].mlp.c_fc.weight.device |
| 325 | |
| 326 | def forward(self, x: torch.Tensor, attn_mask: Optional[torch.Tensor] = None): |
| 327 | for r in self.resblocks: |
| 328 | x = r(x, attn_mask=attn_mask) |
| 329 | return x |
| 330 | |
| 331 | |
| 332 | class VisionTransformer(nn.Module): |