Spaces:
Build error
Build error
File size: 3,608 Bytes
3346833 | 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 87 | import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange
from einops.layers.torch import Rearrange
import random,math
class shift_patch(nn.Module):
def __init__(self,patch_size=16):
super().__init__()
self.patch_size = patch_size // 2
def forward(self, x):
bs,_, _, chanel = x.size()
shifts = ((-self.patch_size, self.patch_size, -self.patch_size, self.patch_size),
(self.patch_size, -self.patch_size, -self.patch_size, self.patch_size),
(-self.patch_size, self.patch_size, self.patch_size, -self.patch_size),
(self.patch_size, -self.patch_size, self.patch_size, -self.patch_size))
shifted_x = list(map(lambda shift: F.pad(x, shift), shifts))
return shifted_x
def posemb_sincos_2d(h, w, dim, temperature: int = 10000, dtype = torch.float32):
y, x = torch.meshgrid(torch.arange(h), torch.arange(w), indexing="ij")
assert (dim % 4) == 0, "feature dimension must be multiple of 4 for sincos emb"
omega = torch.arange(dim // 4) / (dim // 4 - 1)
omega = 1.0 / (temperature ** omega)
y = y.flatten()[:, None] * omega[None, :]
x = x.flatten()[:, None] * omega[None, :]
pe = torch.cat((x.sin(), x.cos(), y.sin(), y.cos()), dim=1)
return pe.type(dtype)
class PositionalEncoding(nn.Module):
def __init__(self, d_model, max_seq_length):
super(PositionalEncoding, self).__init__()
pe = torch.zeros(max_seq_length, d_model)
position = torch.arange(0, max_seq_length, dtype=torch.float).unsqueeze(1)
div_term = torch.exp(torch.arange(0, d_model, 2).float() * -(math.log(10000.0) / d_model))
pe[:, 0::2] = torch.sin(position * div_term)
pe[:, 1::2] = torch.cos(position * div_term)
self.register_buffer('pe', pe.unsqueeze(0))
def forward(self, x):
return x + self.pe[:, :x.size(1)]
class patch_embedding(nn.Module):
def __init__(self,height=224,width=224,n_channels=3,patch_size=16,dim=512):
super().__init__()
assert height%patch_size==0 and width%patch_size==0 ,"Height and Width should be multiples of patch size wich is {0}".format(patch_size)
#self.class_token = nn.Parameter(torch.randn(1, 1, dim))
self.pos_embedding = posemb_sincos_2d(
h = height // patch_size,
w = width // patch_size,
dim = dim,
)
self.patch_size = patch_size
self.n_patchs = height*width//(patch_size**2)
#self.embedding = torch.nn.Parameter(torch.randn(1,self.n_patchs+1, dim))
self.projection = nn.Sequential(nn.LayerNorm(patch_size*patch_size*n_channels*5),
nn.Linear(patch_size*patch_size*n_channels*5,dim),
nn.LayerNorm(dim)
)
self.shift = shift_patch(patch_size=patch_size)
self.patch = nn.Sequential(
Rearrange("b c (h p1) (w p2) -> b (h w) (p1 p2 c)", p1 = patch_size, p2 = patch_size),
)
def forward(self, x):
#x bs,h,w,c
embedding = self.pos_embedding.to(x.device)
left_up,right_up,left_down,right_down = self.shift(x)
x = torch.cat([x,left_up,right_up,left_down,right_down],dim=1)
#first we resize the inputs Bs,Num_patchs,*
x = self.patch(x)
#projection on the dim of model
x = self.projection(x)
outputs = x + embedding
return outputs
|