File size: 2,853 Bytes
88ae77c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 |
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
|