The autoencoder model to enhance images in an image to image translation fashion. This code is built on top of the vit-pytorch code https://github.com/lucidrains/vit-pytorch. Args: encoder (model): the defined encoder, hete it is a ViT decoder_dim (int): decoder dim (em
| 28 | logger = logging.getLogger('base') |
| 29 | |
| 30 | class BinModel(nn.Module): |
| 31 | """ |
| 32 | The autoencoder model to enhance images in an image to image translation fashion. |
| 33 | This code is built on top of the vit-pytorch code https://github.com/lucidrains/vit-pytorch. |
| 34 | |
| 35 | Args: |
| 36 | encoder (model): the defined encoder, hete it is a ViT |
| 37 | decoder_dim (int): decoder dim (embedding size) |
| 38 | decoder_depth (int): number of decoder layers |
| 39 | decoder_heads (int): number of decoder heads |
| 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 | |
| 67 | # get patches and their number |
| 68 | patches = self.to_patch(img) |
| 69 | print('patch shape:', patches.shape) |
| 70 | _, num_patches, *_ = patches.shape |
| 71 | |
| 72 | # project pixel patches to tokens and add positions |
| 73 | tokens = self.patch_to_emb(patches) |
| 74 | tokens = tokens + self.encoder.pos_embedding[:, 1:(num_patches + 1)] |
| 75 | |
| 76 | # encode tokens by the encoder |
| 77 | encoded_tokens = self.encoder.transformer(tokens) |
| 78 | |
| 79 | # project encoder to decoder dimensions, if they are not equal. |
| 80 | decoder_tokens = self.enc_to_dec(encoded_tokens) |
| 81 | |
| 82 | # decode tokens with decoder |
| 83 | decoded_tokens = self.decoder(decoder_tokens) |
| 84 | |
| 85 | # project tokens to pixels |
| 86 | pred_pixel_values = self.to_pixels(decoded_tokens) |
| 87 |