| 26 | % (type(self).__name__, num_params / 1000000)) |
| 27 | |
| 28 | def init_weights(self, init_type='normal', gain=0.02): |
| 29 | def init_func(m): |
| 30 | classname = m.__class__.__name__ |
| 31 | if classname.find('BatchNorm2d') != -1: |
| 32 | if hasattr(m, 'weight') and m.weight is not None: |
| 33 | init.normal_(m.weight.data, 1.0, gain) |
| 34 | if hasattr(m, 'bias') and m.bias is not None: |
| 35 | init.constant_(m.bias.data, 0.0) |
| 36 | elif hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1): |
| 37 | if init_type == 'normal': |
| 38 | init.normal_(m.weight.data, 0.0, gain) |
| 39 | elif init_type == 'xavier': |
| 40 | init.xavier_normal_(m.weight.data, gain=gain) |
| 41 | elif init_type == 'xavier_uniform': |
| 42 | init.xavier_uniform_(m.weight.data, gain=1.0) |
| 43 | elif init_type == 'kaiming': |
| 44 | init.kaiming_normal_(m.weight.data, a=0, mode='fan_in') |
| 45 | elif init_type == 'orthogonal': |
| 46 | init.orthogonal_(m.weight.data, gain=gain) |
| 47 | elif init_type == 'none': # uses pytorch's default init method |
| 48 | m.reset_parameters() |
| 49 | else: |
| 50 | raise NotImplementedError('initialization method [%s] is not implemented' % init_type) |
| 51 | if hasattr(m, 'bias') and m.bias is not None: |
| 52 | init.constant_(m.bias.data, 0.0) |
| 53 | |
| 54 | self.apply(init_func) |
| 55 | |
| 56 | # propagate to children |
| 57 | for m in self.children(): |
| 58 | if hasattr(m, 'init_weights'): |
| 59 | m.init_weights(init_type, gain) |