| 74 | |
| 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) |
| 110 | |
| 111 | x = self.proj_in(x) # (b, 512, 1024) |
| 112 | |
| 113 | for attn, ff in self.layers: |
| 114 | latents = attn(x, latents) + latents # b 16 1024 |
| 115 | latents = ff(latents) + latents |
| 116 | |
| 117 | latents = self.proj_out(latents) |
| 118 | return self.norm_out(latents) |
| 119 | |
| 120 | |
| 121 | class MultiProjModel(nn.Module): |