| 40 | decoder_dim_head (int): decoder head dimension |
| 41 | """ |
| 42 | def __init__( |
| 43 | self, |
| 44 | *, |
| 45 | encoder, |
| 46 | decoder_dim, |
| 47 | decoder_depth = 1, |
| 48 | decoder_heads = 8, |
| 49 | decoder_dim_head = 64 |
| 50 | ): |
| 51 | super().__init__() |
| 52 | # extract hyperparameters and functions from the ViT encoder. |
| 53 | self.encoder = encoder |
| 54 | num_patches, encoder_dim = encoder.pos_embedding.shape[-2:] |
| 55 | self.to_patch, self.patch_to_emb = encoder.to_patch_embedding[:2] |
| 56 | pixel_values_per_patch = self.patch_to_emb.weight.shape[-1] |
| 57 | |
| 58 | # define your decoder here |
| 59 | self.enc_to_dec = nn.Linear(encoder_dim, decoder_dim) if encoder_dim != decoder_dim else nn.Identity() |
| 60 | self.mask_token = nn.Parameter(torch.randn(decoder_dim)) |
| 61 | self.decoder = Transformer(dim = decoder_dim, depth = decoder_depth, heads = decoder_heads, dim_head = decoder_dim_head, mlp_dim = decoder_dim * 4) |
| 62 | self.decoder_pos_emb = nn.Embedding(num_patches, decoder_dim) |
| 63 | self.to_pixels = nn.Linear(decoder_dim, pixel_values_per_patch) |
| 64 | |
| 65 | def forward(self, img): |
| 66 | |