File size: 1,536 Bytes
af2b273 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 | from torch import nn
from dynamic_network_architectures.building_blocks.residual import BasicBlockD, BottleneckD
class InitWeights_He(object):
def __init__(self, neg_slope: float = 1e-2):
self.neg_slope = neg_slope
def __call__(self, module):
if isinstance(module, nn.Conv3d) or isinstance(module, nn.Conv2d) or isinstance(module, nn.ConvTranspose2d) or isinstance(module, nn.ConvTranspose3d):
module.weight = nn.init.kaiming_normal_(module.weight, a=self.neg_slope)
if module.bias is not None:
module.bias = nn.init.constant_(module.bias, 0)
class InitWeights_XavierUniform(object):
def __init__(self, gain: int = 1):
self.gain = gain
def __call__(self, module):
if isinstance(module, nn.Conv3d) or isinstance(module, nn.Conv2d) or isinstance(module, nn.ConvTranspose2d) or isinstance(module, nn.ConvTranspose3d):
module.weight = nn.init.xavier_uniform_(module.weight, self.gain)
if module.bias is not None:
module.bias = nn.init.constant_(module.bias, 0)
def init_last_bn_before_add_to_0(module):
if isinstance(module, BasicBlockD):
module.conv2.norm.weight = nn.init.constant_(module.conv2.norm.weight, 0)
module.conv2.norm.bias = nn.init.constant_(module.conv2.norm.bias, 0)
if isinstance(module, BottleneckD):
module.conv3.norm.weight = nn.init.constant_(module.conv3.norm.weight, 0)
module.conv3.norm.bias = nn.init.constant_(module.conv3.norm.bias, 0)
|