| 580 | return xi, xt |
| 581 | |
| 582 | def init_weights(self): |
| 583 | # embeddings |
| 584 | nn.init.normal_(self.textual.token_embedding.weight, std=0.02) |
| 585 | nn.init.normal_(self.visual.patch_embedding.weight, std=0.1) |
| 586 | |
| 587 | # attentions |
| 588 | for modality in ['visual', 'textual']: |
| 589 | dim = self.vision_dim if modality == 'visual' else self.text_dim |
| 590 | transformer = getattr(self, modality).transformer |
| 591 | proj_gain = (1.0 / math.sqrt(dim)) * ( |
| 592 | 1.0 / math.sqrt(2 * len(transformer))) |
| 593 | attn_gain = 1.0 / math.sqrt(dim) |
| 594 | mlp_gain = 1.0 / math.sqrt(2.0 * dim) |
| 595 | for block in transformer: |
| 596 | nn.init.normal_(block.attn.to_qkv.weight, std=attn_gain) |
| 597 | nn.init.normal_(block.attn.proj.weight, std=proj_gain) |
| 598 | nn.init.normal_(block.mlp[0].weight, std=mlp_gain) |
| 599 | nn.init.normal_(block.mlp[2].weight, std=proj_gain) |
| 600 | |
| 601 | def param_groups(self): |
| 602 | groups = [{ |