Initialize weights with He. initialization and zero out the biases
(self)
| 42 | self.scale = nn.Parameter(torch.ones(1) * 2) |
| 43 | |
| 44 | def _initialize_weights(self): |
| 45 | """Initialize weights with He. initialization and zero out the biases |
| 46 | """ |
| 47 | for m in self.conv_layers: |
| 48 | if isinstance(m, nn.Conv2d): |
| 49 | n = m.kernel_size[0] * m.kernel_size[1] * m.in_channels |
| 50 | init.normal_(m.weight, mean=0.0, std=np.sqrt(2. / n)) |
| 51 | if m.bias is not None: |
| 52 | init.zeros_(m.bias) |
| 53 | init.zeros_(self.final_proj.weight) |
| 54 | if self.final_proj.bias is not None: |
| 55 | init.zeros_(self.final_proj.bias) |
| 56 | |
| 57 | def forward(self, x): |
| 58 | if x.ndim == 5: |