| import torch.nn as nn |
|
|
|
|
| def init_weights(model: nn.Module, conv_std_or_gain: float = 0.02, other_std: float = 0.02): |
| """ |
| :param model: the model to be inited |
| :param conv_std_or_gain: how to init every conv layer `m` |
| > 0: nn.init.trunc_normal_(m.weight.data, std=conv_std_or_gain) |
| < 0: nn.init.xavier_normal_(m.weight.data, gain=-conv_std_or_gain) |
| :param other_std: how to init every linear layer or embedding layer |
| use nn.init.trunc_normal_(m.weight.data, std=other_std) |
| """ |
| skip = abs(conv_std_or_gain) > 10 |
| if skip: return |
| print(f'[init_weights] {type(model).__name__} with {"std" if conv_std_or_gain > 0 else "gain"}={abs(conv_std_or_gain):g}') |
| for m in model.modules(): |
| if isinstance(m, nn.Linear): |
| nn.init.trunc_normal_(m.weight.data, std=other_std) |
| if m.bias is not None: |
| nn.init.constant_(m.bias.data, 0.) |
| elif isinstance(m, nn.Embedding): |
| nn.init.trunc_normal_(m.weight.data, std=other_std) |
| if m.padding_idx is not None: |
| m.weight.data[m.padding_idx].zero_() |
| elif isinstance(m, (nn.Conv1d, nn.Conv2d, nn.ConvTranspose1d, nn.ConvTranspose2d)): |
| nn.init.trunc_normal_(m.weight.data, std=conv_std_or_gain) if conv_std_or_gain > 0 else nn.init.xavier_normal_(m.weight.data, gain=-conv_std_or_gain) |
| if hasattr(m, 'bias') and m.bias is not None: |
| nn.init.constant_(m.bias.data, 0.) |
| elif isinstance(m, (nn.LayerNorm, nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d, nn.SyncBatchNorm, nn.GroupNorm, nn.InstanceNorm1d, nn.InstanceNorm2d, nn.InstanceNorm3d)): |
| if m.bias is not None: |
| nn.init.constant_(m.bias.data, 0.) |
| if m.weight is not None: |
| nn.init.constant_(m.weight.data, 1.) |
|
|
