Spaces:
Build error
Build error
| 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 | |