Spaces:
Sleeping
Sleeping
| import copy | |
| import os | |
| import numpy as np | |
| import pytorch_lightning as pl | |
| import torch | |
| import torch.nn as nn | |
| from einops import rearrange | |
| from torchmetrics.functional import accuracy | |
| from torchmetrics.functional.classification import multiclass_recall, multiclass_precision | |
| from x_transformers import Encoder, Decoder | |
| ON_EPOCH = True | |
| ON_STEP = False | |
| BATCH_SIZE = 64 | |
| TARGET_SIZE = (64, 64) | |
| SPLIT_RATE = 0.8 | |
| ROOT_DIR_DATA = "/kaggle/input/ant-data-new/data" | |
| class PatchEmbed(nn.Module): | |
| """Image to Patch Embedding""" | |
| def __init__(self, img_size=TARGET_SIZE[0], patch_size=4, in_chans=3, embed_dim=64): | |
| super().__init__() | |
| if isinstance(img_size, int): | |
| img_size = img_size, img_size | |
| if isinstance(patch_size, int): | |
| patch_size = patch_size, patch_size | |
| # calculate the number of patches | |
| self.patch_shape = (img_size[0] // patch_size[0], img_size[1] // patch_size[1]) | |
| # convolutional layer to convert the image into patches | |
| self.conv = nn.Conv2d( | |
| in_chans, embed_dim, kernel_size=patch_size, stride=patch_size | |
| ) | |
| def forward(self, x): | |
| x = self.conv(x) | |
| # flatten the patches | |
| x = rearrange(x, 'b e h w -> b (h w) e') | |
| return x | |
| class ViTIJEPA(nn.Module): | |
| def __init__(self, img_size, patch_size, in_chans, embed_dim, enc_depth, num_heads, | |
| num_classes, post_emb_norm=False, | |
| layer_dropout=0.): | |
| super().__init__() | |
| self.layer_dropout = layer_dropout | |
| self.patch_embed = PatchEmbed(img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim) | |
| self.num_tokens = self.patch_embed.patch_shape[0] * self.patch_embed.patch_shape[1] | |
| self.pos_embedding = nn.Parameter(torch.randn(1, self.num_tokens, embed_dim)) | |
| self.post_emb_norm = nn.LayerNorm(embed_dim) if post_emb_norm else nn.Identity() | |
| self.student_encoder = Encoder( | |
| dim=embed_dim, | |
| heads=num_heads, | |
| depth=enc_depth, | |
| layer_dropout=self.layer_dropout, | |
| flash=True | |
| ) | |
| self.average_pool = nn.AvgPool1d((embed_dim), stride=1) | |
| # mlp head | |
| self.mlp_head = nn.Sequential( | |
| nn.LayerNorm(self.num_tokens), | |
| nn.Linear(self.num_tokens, num_classes), | |
| ) | |
| def forward(self, x): | |
| x = self.patch_embed(x) | |
| b, n, e = x.shape | |
| # add the positional embeddings | |
| x = x + self.pos_embedding | |
| # normalize the embeddings | |
| x = self.post_emb_norm(x) | |
| # if mode is test, we get return full embedding: | |
| x = self.student_encoder(x) | |
| x = self.average_pool(x) # conduct average pool like in paper | |
| x = x.squeeze(-1) | |
| x = self.mlp_head(x) # pass through mlp head | |
| return x | |