| import torch | |
| import torchvision | |
| from torch import nn | |
| from torchvision import datasets, transforms | |
| PATCH_SIZE = 16 | |
| class PatchEmbeddings(nn.Module): | |
| def __init__(self, in_channels: int=3, | |
| patch_size: int=16, | |
| embedding_dim: int=768): | |
| super().__init__() | |
| self.generate_patches = nn.Conv2d(in_channels=in_channels, | |
| out_channels=embedding_dim, | |
| kernel_size=patch_size, | |
| stride=patch_size, padding=0) | |
| self.flatten = nn.Flatten(start_dim=2, end_dim=3) | |
| def forward(self, x: torch.Tensor): | |
| image_resolution = x.shape[-1] | |
| assert image_resolution % PATCH_SIZE == 0, f"Image size must be divisible by patch size!" | |
| return self.flatten(self.generate_patches(x)).permute(0, 2, 1) | |
| class MultiheadSelfAttention(nn.Module): | |
| def __init__(self, embedding_dim: int=768, | |
| num_heads: int=12, attn_dropout: int=0): | |
| super().__init__() | |
| self.layer_norm = nn.LayerNorm(normalized_shape=embedding_dim) | |
| self.multihead_attn = nn.MultiheadAttention(embed_dim=embedding_dim, num_heads=num_heads, | |
| dropout=attn_dropout, batch_first=True) | |
| def forward(self, x: torch.Tensor): | |
| x = self.layer_norm(x) | |
| attn_output, _ = self.multihead_attn(query=x, key=x, value=x, | |
| need_weights=False) | |
| return attn_output | |
| class MLPBlock(nn.Module): | |
| def __init__(self, embedding_dim: int=768, | |
| mlp_size: int=3072, dropout: int=0.1): | |
| super().__init__() | |
| self.layer_norm = nn.LayerNorm(normalized_shape=embedding_dim) | |
| self.mlp = nn.Sequential( | |
| nn.Linear(in_features=embedding_dim, | |
| out_features=mlp_size), | |
| nn.GELU(), | |
| nn.Dropout(p=dropout), | |
| nn.Linear(in_features=mlp_size, | |
| out_features=embedding_dim), | |
| nn.Dropout(p=dropout) | |
| ) | |
| def forward(self, x: torch.Tensor): | |
| return self.mlp(self.layer_norm(x)) | |
| class TransformerEncoderBlock(nn.Module): | |
| def __init__(self, embedding_dim: int=768, | |
| mlp_size: int=3072, num_heads: int=12, | |
| mlp_dropout: int=0.1, attn_dropout: int=0): | |
| super().__init__() | |
| self.msa_block = MultiheadSelfAttention(embedding_dim=embedding_dim, | |
| num_heads=num_heads, attn_dropout=attn_dropout) | |
| self.mlp_block = MLPBlock(embedding_dim=embedding_dim, | |
| mlp_size=mlp_size, dropout=mlp_dropout) | |
| def forward(self, x: torch.Tensor): | |
| x = self.msa_block(x) + x | |
| x = self.mlp_block(x) + x | |
| return x | |
| class VisionTransformer(nn.Module): | |
| def __init__(self, img_size: int=IMG_SIZE, | |
| in_channels: int=3, patch_size: int=16, | |
| num_transformer_layers: int=12, embedding_dim: int=768, | |
| mlp_size: int=3072, num_heads: int=12, | |
| attn_dropout: int=0, mlp_dropout: int=0.1, | |
| embedding_dropout: int=0.1, num_classes: int=38): | |
| super().__init__() | |
| assert img_size % patch_size == 0, f"Image size must be divisible by patch size!" | |
| self.num_patches = (img_size * img_size) // patch_size**2 | |
| self.class_embedding = nn.Parameter(data=torch.randn(1, 1, embedding_dim), | |
| requires_grad=True) | |
| self.position_embedding = nn.Parameter(data=torch.randn(1, self.num_patches+1, embedding_dim), | |
| requires_grad=True) | |
| self.embedding_dropout = nn.Dropout(p=embedding_dropout) | |
| self.patch_embeddings = PatchEmbeddings(in_channels=in_channels, | |
| patch_size=patch_size, embedding_dim=embedding_dim) | |
| self.transformer_encoder = nn.Sequential(*[TransformerEncoderBlock(embedding_dim=embedding_dim, | |
| num_heads=num_heads, mlp_size=mlp_size, | |
| mlp_dropout=mlp_dropout) for _ in range(num_transformer_layers)]) | |
| self.classifier = nn.Sequential( | |
| nn.LayerNorm(normalized_shape=embedding_dim), | |
| nn.Linear(in_features=embedding_dim, | |
| out_features=num_classes) | |
| ) | |
| def forward(self, x: torch.Tensor): | |
| batch_size = x.shape[0] | |
| class_token = self.class_embedding.expand(batch_size, -1, -1) | |
| x = self.patch_embeddings(x) | |
| x = torch.cat((class_token, x), dim=1) | |
| x = self.position_embedding + x | |
| x = self.embedding_dropout(x) | |
| x = self.transformer_encoder(x) | |
| x = self.classifier(x[:, 0]) | |
| return x | |
| with open("class_names.ob", "rb") as fp: | |
| class_names = pickle.load(fp) | |
| vision_transformer = VisionTransformer(num_classes=len(class_names)) | |