(self)
| 397 | self.init_weights() |
| 398 | |
| 399 | def init_weights(self) -> None: |
| 400 | std = 1.0 / math.sqrt(self._query_dim) |
| 401 | torch.nn.init.trunc_normal_(self.q_proj.weight, std=std, a=-3 * std, b=3 * std) |
| 402 | std = 1.0 / math.sqrt(self._context_dim) |
| 403 | torch.nn.init.trunc_normal_(self.k_proj.weight, std=std, a=-3 * std, b=3 * std) |
| 404 | torch.nn.init.trunc_normal_(self.v_proj.weight, std=std, a=-3 * std, b=3 * std) |
| 405 | |
| 406 | std = 1.0 / math.sqrt(self._inner_dim) |
| 407 | torch.nn.init.trunc_normal_(self.output_proj.weight, std=std, a=-3 * std, b=3 * std) |
| 408 | |
| 409 | for layer in self.q_norm, self.k_norm, self.v_norm: |
| 410 | if hasattr(layer, "reset_parameters"): |
| 411 | layer.reset_parameters() |
| 412 | |
| 413 | def compute_qkv( |
| 414 | self, |
no test coverage detected