| 330 | |
| 331 | |
| 332 | class VisionTransformer(nn.Module): |
| 333 | |
| 334 | def __init__( |
| 335 | self, |
| 336 | image_size: int, |
| 337 | patch_size: int, |
| 338 | width: int, |
| 339 | layers: int, |
| 340 | heads: int, |
| 341 | mlp_ratio: float, |
| 342 | n_queries: int = 256, |
| 343 | output_dim: int = 512, |
| 344 | **kwargs |
| 345 | ): |
| 346 | super().__init__() |
| 347 | image_height, image_width = self.image_size = (image_size, image_size) |
| 348 | patch_height, patch_width = self.patch_size = (patch_size, patch_size) |
| 349 | self.grid_size = (image_height // patch_height, image_width // patch_width) |
| 350 | self.output_dim = output_dim |
| 351 | |
| 352 | mean = (0.48145466, 0.4578275, 0.40821073) |
| 353 | std = (0.26862954, 0.26130258, 0.27577711) |
| 354 | self.image_transform = transforms.Compose([ |
| 355 | transforms.Resize( |
| 356 | (image_size, image_size), |
| 357 | interpolation=InterpolationMode.BICUBIC |
| 358 | ), |
| 359 | transforms.ToTensor(), |
| 360 | transforms.Normalize(mean=mean, std=std), |
| 361 | ]) |
| 362 | |
| 363 | self.conv1 = nn.Conv2d(in_channels=3, out_channels=width, kernel_size=patch_size, stride=patch_size, bias=False) |
| 364 | |
| 365 | # class embeddings and positional embeddings |
| 366 | scale = width ** -0.5 |
| 367 | self.positional_embedding = nn.Parameter(scale * torch.randn(256, width)) |
| 368 | |
| 369 | norm_layer = partial(nn.LayerNorm, eps=1e-6) |
| 370 | act_layer = nn.GELU |
| 371 | |
| 372 | self.ln_pre = norm_layer(width) |
| 373 | self.transformer = TransformerBlock( |
| 374 | width, |
| 375 | layers, |
| 376 | heads, |
| 377 | mlp_ratio, |
| 378 | act_layer=act_layer, |
| 379 | norm_layer=norm_layer, |
| 380 | ) |
| 381 | |
| 382 | self.attn_pool = Resampler( |
| 383 | grid_size=int(math.sqrt(n_queries)), |
| 384 | embed_dim=output_dim, |
| 385 | num_heads=output_dim // 128, |
| 386 | kv_dim=width, |
| 387 | norm_layer=norm_layer, |
| 388 | ) |
| 389 | self.ln_post = norm_layer(output_dim) |