ZF/vit-pytorch
1
0
Fork 0
mirror of https://github.com/lucidrains/vit-pytorch.git synced 2025-12-30 08:02:29 +00:00
Code Issues Packages Projects Releases Wiki Activity

go for multi-headed rmsnorm for the qknorm on hypersphere vit

Browse Source
This commit is contained in:
Phil Wang
2024-10-10 08:09:58 -07:00
parent f50d7d1436
commit 74b62009f8
2 changed files with 6 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
setup.py
View File

@@ -6,7 +6,7 @@ with open('README.md') as f:
setup( setup(
name = 'vit-pytorch', name = 'vit-pytorch',
packages = find_packages(exclude=['examples']), packages = find_packages(exclude=['examples']),
version = '1.8.1', version = '1.8.2',
license='MIT', license='MIT',
description = 'Vision Transformer (ViT) - Pytorch', description = 'Vision Transformer (ViT) - Pytorch',
long_description=long_description, long_description=long_description,

7
vit_pytorch/normalized_vit.py
View File

@@ -76,7 +76,8 @@ class Attention(Module):
self.dropout = dropout self.dropout = dropout
self.qk_scale = nn.Parameter(torch.ones(dim_head) * (dim_head ** 0.25)) self.q_scale = nn.Parameter(torch.ones(dim_inner) * (dim_head ** 0.25))
self.k_scale = nn.Parameter(torch.ones(dim_inner) * (dim_head ** 0.25))
self.split_heads = Rearrange('b n (h d) -> b h n d', h = heads) self.split_heads = Rearrange('b n (h d) -> b h n d', h = heads)
self.merge_heads = Rearrange('b h n d -> b n (h d)') self.merge_heads = Rearrange('b h n d -> b n (h d)')
@@ -89,12 +90,14 @@ class Attention(Module):
): ):
q, k, v = self.to_q(x), self.to_k(x), self.to_v(x) q, k, v = self.to_q(x), self.to_k(x), self.to_v(x)
q = q * self.q_scale
k = k * self.k_scale
q, k, v = map(self.split_heads, (q, k, v)) q, k, v = map(self.split_heads, (q, k, v))
# query key rmsnorm # query key rmsnorm
q, k = map(l2norm, (q, k)) q, k = map(l2norm, (q, k))
q, k = (q * self.qk_scale), (k * self.qk_scale)
# scale is 1., as scaling factor is moved to s_qk (dk ^ 0.25) - eq. 16 # scale is 1., as scaling factor is moved to s_qk (dk ^ 0.25) - eq. 16