Initialize the weights of the model. - Xavier uniform initialization for linear layers - Normal initialization for embeddings - Xavier uniform initialization for parameters
(model: nn.Module)
| 197 | |
| 198 | |
| 199 | def init_weights(model: nn.Module): |
| 200 | """ |
| 201 | Initialize the weights of the model. |
| 202 | - Xavier uniform initialization for linear layers |
| 203 | - Normal initialization for embeddings |
| 204 | - Xavier uniform initialization for parameters |
| 205 | """ |
| 206 | |
| 207 | def _init_weights(m): |
| 208 | if isinstance(m, nn.Linear): |
| 209 | nn.init.xavier_uniform_(m.weight) |
| 210 | if m.bias is not None: |
| 211 | nn.init.zeros_(m.bias) |
| 212 | elif isinstance(m, nn.Embedding): |
| 213 | nn.init.normal_(m.weight, mean=0.0, std=0.02) |
| 214 | elif isinstance(m, nn.Parameter): |
| 215 | nn.init.xavier_uniform_(m.data) |
| 216 | |
| 217 | model.apply(_init_weights) |
| 218 | |
| 219 | # Special handling for audio_head because it's nn.Parameter directly |
| 220 | nn.init.xavier_uniform_(model.audio_head) |
| 221 | |
| 222 | return model |
| 223 | |
| 224 | |
| 225 | def load_model( |
no outgoing calls
no test coverage detected