| 393 | self.text_projection = nn.Parameter(torch.empty(transformer_width, embed_dim)) |
| 394 | |
| 395 | def init_weights(self, pretrained=None): |
| 396 | pretrained = pretrained or self.pretrained |
| 397 | if isinstance(pretrained, str): |
| 398 | checkpoint = torch.jit.load(pretrained, map_location='cpu').float().state_dict() |
| 399 | |
| 400 | state_dict = {} |
| 401 | |
| 402 | for k in checkpoint.keys(): |
| 403 | if k.startswith('transformer.'): |
| 404 | state_dict[k] = checkpoint[k] |
| 405 | |
| 406 | if k == 'positional_embedding' or k == 'text_projection' or k.startswith('token_embedding') or k.startswith('ln_final'): |
| 407 | if k == 'positional_embedding' and checkpoint[k].size(0) > self.context_length: |
| 408 | checkpoint[k] = checkpoint[k][:self.context_length] |
| 409 | print('positional_embedding is tuncated from 77 to', self.context_length) |
| 410 | state_dict[k] = checkpoint[k] |
| 411 | |
| 412 | u, w = self.load_state_dict(state_dict, False) |
| 413 | print(u, w, 'are misaligned params in text encoder') |
| 414 | |
| 415 | |
| 416 | def build_attention_mask(self): |