Skip to content

Commit

Permalink
make nest resilient to dimension that are not divisible by number of …
Browse files Browse the repository at this point in the history
…heads
  • Loading branch information
lucidrains committed May 28, 2021
1 parent daf3abb commit 17cb897
Show file tree
Hide file tree
Showing 2 changed files with 7 additions and 6 deletions.
2 changes: 1 addition & 1 deletion setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
setup(
name = 'vit-pytorch',
packages = find_packages(exclude=['examples']),
version = '0.19.0',
version = '0.19.1',
license='MIT',
description = 'Vision Transformer (ViT) - Pytorch',
author = 'Phil Wang',
Expand Down
11 changes: 6 additions & 5 deletions vit_pytorch/nest.py
Original file line number Diff line number Diff line change
Expand Up @@ -48,16 +48,16 @@ def forward(self, x):
class Attention(nn.Module):
def __init__(self, dim, heads = 8, dropout = 0.):
super().__init__()
assert (dim % heads) == 0, 'dimension must be divisible by number of heads'
dim_head = dim // heads
inner_dim = dim_head * heads
self.heads = heads
self.scale = dim_head ** -0.5

self.attend = nn.Softmax(dim = -1)
self.to_qkv = nn.Conv2d(dim, dim * 3, 1, bias = False)
self.to_qkv = nn.Conv2d(dim, inner_dim * 3, 1, bias = False)

self.to_out = nn.Sequential(
nn.Conv2d(dim, dim, 1),
nn.Conv2d(inner_dim, dim, 1),
nn.Dropout(dropout)
)

Expand Down Expand Up @@ -129,7 +129,8 @@ def __init__(
blocks = 2 ** (num_heirarchies - 1)

seq_len = (fmap_size // blocks) ** 2 # sequence length is held constant across heirarchy
mults = [2 ** i for i in reversed(range(num_heirarchies))]
heirarchies = list(reversed(range(num_heirarchies)))
mults = [2 ** i for i in heirarchies]

layer_heads = list(map(lambda t: t * heads, mults))
layer_dims = list(map(lambda t: t * dim, mults))
Expand All @@ -146,7 +147,7 @@ def __init__(

self.layers = nn.ModuleList([])

for level, heads, (dim_in, dim_out), block_repeat in zip(reversed(range(num_heirarchies)), layer_heads, dim_pairs, block_repeats):
for level, heads, (dim_in, dim_out), block_repeat in zip(heirarchies, layer_heads, dim_pairs, block_repeats):
is_last = level == 0
depth = block_repeat

Expand Down

0 comments on commit 17cb897

Please sign in to comment.