from .utils.network_blocks import * from .utils.network_blocks_2d import * class BaseNetwork(nn.Module): def __init__(self, conv_type): super(BaseNetwork, self).__init__() self.conv_type = conv_type if conv_type == 'gated': self.ConvBlock = GatedConv self.DeconvBlock = GatedDeconv self.ConvBlock2d = GatedConv2d self.DeconvBlock2d = GatedDeconv2d if conv_type == 'partial': self.ConvBlock = PartialConv self.DeconvBlock = PartialDeconv self.ConvBlock2d = PartialConv2d self.DeconvBlock2d = PartialDeconv2d if conv_type == 'vanilla': self.ConvBlock = VanillaConv self.DeconvBlock = VanillaDeconv self.ConvBlock2d = VanillaConv2d self.DeconvBlock2d = VanillaDeconv2d def init_weights(self, init_type='kaiming', gain=0.02): ''' initialize network's weights init_type: normal | xavier | kaiming | orthogonal https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/9451e70673400885567d08a9e97ade2524c700d0/models/networks.py#L39 ''' def init_func(m): classname = m.__class__.__name__ if hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1): if init_type == 'normal': nn.init.normal_(m.weight.data, 0.0, gain) elif init_type == 'xavier': nn.init.xavier_normal_(m.weight.data, gain=gain) elif init_type == 'kaiming': nn.init.kaiming_normal_(m.weight.data, a=0, mode='fan_in') elif init_type == 'orthogonal': nn.init.orthogonal_(m.weight.data, gain=gain) if hasattr(m, 'bias') and m.bias is not None: nn.init.constant_(m.bias.data, 0.0) elif classname.find('BatchNorm2d') != -1: nn.init.normal_(m.weight.data, 1.0, gain) nn.init.constant_(m.bias.data, 0.0) self.apply(init_func)