(
self,
dim=1024,
depth=8,
dim_head=64,
heads=16,
num_queries=8,
embedding_dim=768,
output_dim=1024,
ff_mult=4,
)
| 75 | |
| 76 | class Resampler(nn.Module): |
| 77 | def __init__( |
| 78 | self, |
| 79 | dim=1024, |
| 80 | depth=8, |
| 81 | dim_head=64, |
| 82 | heads=16, |
| 83 | num_queries=8, |
| 84 | embedding_dim=768, |
| 85 | output_dim=1024, |
| 86 | ff_mult=4, |
| 87 | ): |
| 88 | super().__init__() |
| 89 | |
| 90 | self.latents = nn.Parameter(torch.randn(1, num_queries, dim) / dim**0.5) |
| 91 | |
| 92 | self.proj_in = nn.Linear(embedding_dim, dim) |
| 93 | |
| 94 | self.proj_out = nn.Linear(dim, output_dim) |
| 95 | self.norm_out = nn.LayerNorm(output_dim) |
| 96 | |
| 97 | self.layers = nn.ModuleList([]) |
| 98 | for _ in range(depth): |
| 99 | self.layers.append( |
| 100 | nn.ModuleList( |
| 101 | [ |
| 102 | PerceiverAttention(dim=dim, dim_head=dim_head, heads=heads), |
| 103 | FeedForward(dim=dim, mult=ff_mult), |
| 104 | ] |
| 105 | ) |
| 106 | ) |
| 107 | |
| 108 | def forward(self, x): # x (b, 512, 1) |
| 109 | latents = self.latents.repeat(x.size(0), 1, 1) |
no test coverage detected