So, i have tried to implement Vision Transformers from scratch, but somehow the layer to_qvk and the W_0 keep initialized on None, even though theres not a single detached layer on this code. Can anyone figured out why? thankyou.
def expand_to_batch(tensor, desired_size):
tile = desired_size // tensor.shape[0]
return repeat(tensor, 'b ... -> (b tile) ...', tile=tile)
def compute_mhsa(q, k, v, scale_factor=1, mask=None):
scaled_dot_prod = torch.einsum('... i d , ... j d -> ... i j', q, k)
attention = torch.softmax(scaled_dot_prod, dim=-1)
return torch.einsum('... i j , ... j d -> ... i d', attention, v)
class MultiHeadSelfAttention(nn.Module):
def __init__(self, dim, heads=4, dim_head=None):
super().__init__()
self.dim_head = (int(dim / heads)) if dim_head is None else dim_head
_dim = self.dim_head * heads
self.heads = heads
self.to_qvk = nn.Linear(dim, _dim * 3, bias=False)
self.W_0 = nn.Linear(_dim, dim, bias=False)
self.rearrange = Rearrange("b h t d -> b t (h d)")
def forward(self, x, mask=None):
logger.info("============ MULTI HEAD SELF ATTENTION ============\n")
assert x.dim() == 3
qkv = self.to_qvk(x)
logger.info(f"Hasil QKV \n{qkv}")
q, k, v = nn.Parameter(rearrange(qkv, 'b t (d k h ) -> k b h t d ', k=3, h=self.heads))
logger.info(f"Hasil : Q \n{q}")
logger.info(f"Hasil : K \n{k}")
logger.info(f"Hasil : V \n{v}")
out = nn.Parameter(compute_mhsa(q, k, v, mask=mask))
logger.info(f"Hasil Scaled dot Product : \n{out}")
out = self.rearrange(out)
logger.info(f"Hasil Concat: \n{out}")
out = self.W_0(out)
logger.info(f"Hasil Linear : \n{out}")
return out
class TransformerBlock(nn.Module):
def __init__(self, dim, heads=8, dim_head=None,
dim_linear_block=1024, activation=nn.GELU,
mhsa=None, prenorm=False):
super().__init__()
self.mhsa = MultiHeadSelfAttention(dim=dim, heads=heads, dim_head=dim_head)
self.prenorm = False
self.norm_1 = nn.LayerNorm(dim)
self.gelu = nn.GELU()
def forward(self, x, mask=None):
logger.info(f"======================== Transformer Block ========================")
logger.info("Before MHSA")
y = self.mhsa(x)
logger.info(f"After MHSA \n{y}")
y = self.norm_1(x)
logger.info(f"After Normalization \n{y}")
out = self.gelu(y)
logger.info(f"After GeLU : \n{out}")
logger.info(f"Hasil Normalization \n {out}")
return out
class TransformerEncoder(nn.Module):
def __init__(self, dim, blocks=6, heads=8, dim_head=None, dim_linear_block=1024, prenorm=False):
super().__init__()
self.transformer_block = TransformerBlock(dim,heads,dim_head,dim_linear_block,prenorm=prenorm)
def forward(self, x, mask=None):
logger.info("======================== Transformer Encoder ========================")
x = self.transformer_block(x,mask)
logger.info(f"Hasil Forward Transformer Encoder {x}")
return x
class ViT(nn.Module):
def __init__(self, *,
img_dim = 4,
in_channels=3,
patch_dim=2,
num_classes=3,
dim=3,
blocks=1,
heads=3,
dim_linear_block=3,
dim_head=None,
transformer=None,
classification=None):
super().__init__()
assert img_dim % patch_dim == 0, f'patch size {patch_dim} not divisible by img dim {img_dim}'
self.num_classes = num_classes
self.p = patch_dim
self.classification = classification
tokens = (img_dim // patch_dim) ** 2
self.token_dim = in_channels * (patch_dim ** 2)
self.dim = dim
self.dim_head = (int(self.dim / heads)) if dim_head is None else dim_head
self.project_patches = nn.Linear(self.token_dim, self.dim,bias=False)
self.cls_token = nn.Parameter(torch.randn(1, 1, self.dim))
self.pos_emb1D = nn.Parameter(torch.randn(tokens + 1, self.dim))
if self.classification:
self.mlp_head = nn.Linear(self.dim, num_classes)
if transformer is None:
self.transformer = TransformerEncoder(self.dim, blocks=blocks, heads=heads,
dim_head=self.dim_head,
dim_linear_block=dim_linear_block)
else:
self.transformer = transformer
self.softmax = nn.Softmax()
def forward(self, img, mask=None):
logger.info("======================== Vision Transformer ========================")
logger.info(f"Image Shape : \n{img.shape}")
logger.info(f"Image Input : \n{img}")
img_patches = rearrange(img,
'b c (patch_x x) (patch_y y) -> b (x y) (patch_x patch_y c)',
patch_x=self.p, patch_y=self.p)
logger.info(f"Image Patches Shape : \n{img_patches.shape}")
logger.info(f"Image Patches : \n{img_patches}")
batch_size, tokens, _ = img_patches.shape
logger.info(f"Batch Size : \n {batch_size}")
logger.info(f"Tokens : \n{tokens}")
img_patches = self.project_patches(img_patches)
logger.info(f"project patches size : \n{img_patches.shape}")
logger.info(f"Project Patches : \n{img_patches}")
img_patches = torch.cat((expand_to_batch(self.cls_token, desired_size=batch_size), img_patches), dim=1)
logger.info(f"project patches after expand to token : \n{img_patches.shape}")
logger.info(f"Images Patches after matrix multiply (cat) : \n{img_patches}")
img_patches = img_patches + self.pos_emb1D[:tokens + 1, :]
logger.info(f"Image patches after Pos embed with tokens : \n{img_patches}")
y = self.transformer(img_patches, mask)
logger.info(f"Image after Transformer : \n{y}")
out = nn.Linear(self.dim, self.num_classes)
out = out(y[:,0,:])
logger.info(f"MLP Head : \n{out}")
return out
and heres the log info
for layer, weight in vit.named_parameters():
print(f"==================Layer {layer}==================================")
if "bias" in layer:
print(f"Bias {weight}")
print(f"Bias Gradient {weight.grad}")
elif "weight" in layer:
print(f" Weight {weight}")
print(f"Gradient {weight.grad}")
==================Layer cls_token==================================
Gradient tensor([[[-4.8729e-05, 1.3597e-05, 3.5132e-05]]])
==================Layer pos_emb1D==================================
Gradient tensor([[-4.8729e-05, 1.3597e-05, 3.5132e-05],
[ 0.0000e+00, 0.0000e+00, 0.0000e+00],
[ 0.0000e+00, 0.0000e+00, 0.0000e+00],
[ 0.0000e+00, 0.0000e+00, 0.0000e+00],
[ 0.0000e+00, 0.0000e+00, 0.0000e+00]])
==================Layer project_patches.weight==================================
Weight Parameter containing:
tensor([[ 0.0490, 0.0490, 0.0490, -0.0480, -0.0480, -0.0480, 0.0490, -0.0490,
0.0450, -0.0490, 0.0480, -0.0490],
[ 0.0490, -0.0480, -0.0480, 0.0480, -0.0480, 0.0470, -0.0490, -0.0490,
-0.0480, 0.0490, -0.0450, -0.0470],
[ 0.0490, -0.0490, 0.0480, -0.0480, -0.0490, 0.0450, -0.0490, -0.0470,
0.0490, -0.0450, -0.0490, -0.0490]], requires_grad=True)
Gradient tensor([[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]])
==================Layer transformer.transformer_block.mhsa.to_qvk.weight==================================
Weight Parameter containing:
tensor([[ 0.0030, 0.0050, 0.0010],
[ 0.0010, -0.0010, -0.0010],
[-0.0030, -0.0030, -0.0030],
[ 0.0030, -0.0030, 0.0030],
[ 0.0040, 0.0040, 0.0040],
[-0.0010, 0.0010, 0.0040],
[-0.0020, -0.0020, -0.0020],
[ 0.0020, -0.0020, 0.0020],
[-0.0020, -0.0020, 0.0010]], requires_grad=True)
Gradient None
==================Layer transformer.transformer_block.mhsa.W_0.weight==================================
Weight Parameter containing:
tensor([[-0.0010, 0.0010, 0.0040],
[-0.0020, 0.0020, -0.0020],
[ 0.0020, -0.0010, -0.0010]], requires_grad=True)
Gradient None
==================Layer transformer.transformer_block.norm_1.weight==================================
Weight Parameter containing:
tensor([-0.0510, 0.0520, 0.0520], requires_grad=True)
Gradient tensor([ 0.1077, -0.1650, -0.2108])
==================Layer transformer.transformer_block.norm_1.bias==================================
Bias Parameter containing:
tensor([-0.0480, -0.0490, 0.0540], requires_grad=True)
Bias Gradient tensor([ 0.3079, 0.1212, -0.2083])
Gradient tensor([ 0.3079, 0.1212, -0.2083])