(self)
| 228 | |
| 229 | |
| 230 | def initialize_weights(self) -> None: |
| 231 | def _basic_init(module): |
| 232 | if isinstance(module, nn.Linear): |
| 233 | torch.nn.init.xavier_uniform_(module.weight) |
| 234 | if module.bias is not None: |
| 235 | nn.init.constant_(module.bias, 0) |
| 236 | |
| 237 | self.apply(_basic_init) |
| 238 | |
| 239 | pos_embed = get_sinusoid_encoding_table(self.num_total_frames, self.hidden_size) |
| 240 | self.pos_embed.data.copy_(pos_embed.unsqueeze(0)) |
| 241 | |
| 242 | w = self.x_embedder.proj.weight.data |
| 243 | nn.init.xavier_uniform_(w.view([w.shape[0], -1])) |
| 244 | nn.init.constant_(self.x_embedder.proj.bias, 0) |
| 245 | |
| 246 | # Initialize timestep embedding MLP: |
| 247 | nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02) |
| 248 | nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02) |
| 249 | |
| 250 | # Zero-out adaLN modulation layers in FMT blocks: |
| 251 | for block in self.blocks: |
| 252 | nn.init.constant_(block.adaLN_modulation[-1].weight, 0) |
| 253 | nn.init.constant_(block.adaLN_modulation[-1].bias, 0) |
| 254 | |
| 255 | # Zero-out output layers: |
| 256 | nn.init.constant_(self.decoder.adaLN_modulation[-1].weight, 0) |
| 257 | nn.init.constant_(self.decoder.adaLN_modulation[-1].bias, 0) |
| 258 | nn.init.constant_(self.decoder.linear.weight, 0) |
| 259 | nn.init.constant_(self.decoder.linear.bias, 0) |
| 260 | |
| 261 | def sequence_embedder(self, sequence, dropout_prob, train=False) -> torch.Tensor: |
| 262 | if train: |
no test coverage detected