| 399 | return self.visual(image.type(self.dtype)) |
| 400 | |
| 401 | def encode_text(self, text): |
| 402 | x = self.token_embedding(text).type(self.dtype) # [batch_size, n_ctx, d_model] |
| 403 | |
| 404 | x = x + self.positional_embedding.type(self.dtype) |
| 405 | x = x.permute(1, 0, 2) # NLD -> LND |
| 406 | x = self.transformer(x) |
| 407 | x = x.permute(1, 0, 2) # LND -> NLD |
| 408 | x = self.ln_final(x).type(self.dtype) |
| 409 | |
| 410 | # x.shape = [batch_size, n_ctx, transformer.width] |
| 411 | # take features from the eot embedding (eot_token is the highest number in each sequence) |
| 412 | x = x[torch.arange(x.shape[0]), text.argmax(dim=-1)] @ self.text_projection |
| 413 | |
| 414 | return x |
| 415 | |
| 416 | def forward(self, image, text): |
| 417 | image_features = self.encode_image(image) |