def __init__(self, startf, maxf, layer_count, latent_size, channels=3): super(EncoderFC, self).__init__() self.maxf = maxf self.startf = startf self.layer_count = layer_count self.channels = channels self.latent_size = latent_size self.fc_1 = ln.Linear(28 * 28, 1024) self.fc_2 = ln.Linear(1024, 1024) self.fc_3 = ln.Linear(1024, latent_size)
def __init__(self, startf=32, maxf=256, layer_count=3, channels=3): super(Discriminator, self).__init__() self.maxf = maxf self.startf = startf self.layer_count = layer_count self.from_rgb = nn.ModuleList() self.channels = channels mul = 2 inputs = startf self.encode_block: nn.ModuleList[DiscriminatorBlock] = nn.ModuleList() resolution = 2**(self.layer_count + 1) for i in range(self.layer_count): outputs = min(self.maxf, startf * mul) self.from_rgb.append(FromRGB(channels, inputs)) fused_scale = resolution >= 128 block = DiscriminatorBlock(inputs, outputs, i == self.layer_count - 1, fused_scale=fused_scale) resolution //= 2 #print("encode_block%d %s" % ((i + 1), millify(count_parameters(block)))) self.encode_block.append(block) inputs = outputs mul *= 2 self.fc2 = ln.Linear(inputs, 1, gain=1)
def __init__(self, inputs, outputs, latent_size, has_first_conv=True, fused_scale=True, layer=0): super(DecodeBlock, self).__init__() self.has_first_conv = has_first_conv self.inputs = inputs self.has_first_conv = has_first_conv self.fused_scale = fused_scale if has_first_conv: if fused_scale: self.conv_1 = ln.ConvTranspose2d(inputs, outputs, 3, 2, 1, bias=False, transform_kernel=True) else: self.conv_1 = ln.Conv2d(inputs, outputs, 3, 1, 1, bias=False) self.blur = Blur(outputs) self.noise_weight_1 = nn.Parameter(torch.Tensor(1, outputs, 1, 1)) self.noise_weight_1.data.zero_() self.bias_1 = nn.Parameter(torch.Tensor(1, outputs, 1, 1)) self.instance_norm_1 = nn.InstanceNorm2d(outputs, affine=False, eps=1e-8) self.style_1 = ln.Linear(latent_size, 2 * outputs, gain=1) self.conv_2 = ln.Conv2d(outputs, outputs, 3, 1, 1, bias=False) self.noise_weight_2 = nn.Parameter(torch.Tensor(1, outputs, 1, 1)) self.noise_weight_2.data.zero_() self.bias_2 = nn.Parameter(torch.Tensor(1, outputs, 1, 1)) self.instance_norm_2 = nn.InstanceNorm2d(outputs, affine=False, eps=1e-8) self.style_2 = ln.Linear(latent_size, 2 * outputs, gain=1) self.layer = layer with torch.no_grad(): self.bias_1.zero_() self.bias_2.zero_()
def __init__(self, startf=32, maxf=256, layer_count=3, latent_size=128, channels=3): super(GeneratorFC, self).__init__() self.maxf = maxf self.startf = startf self.layer_count = layer_count self.channels = channels self.latent_size = latent_size self.fc_1 = ln.Linear(latent_size, 1024) self.fc_2 = ln.Linear(1024, 1024) self.fc_3 = ln.Linear(1024, 28 * 28) self.layer_to_resolution = [28] * 10
def __init__(self, inputs, outputs, latent_size, last=False, fused_scale=True): super(EncodeBlock, self).__init__() self.conv_1 = ln.Conv2d(inputs, inputs, 3, 1, 1, bias=False) # self.conv_1 = ln.Conv2d(inputs + (1 if last else 0), inputs, 3, 1, 1, bias=False) self.bias_1 = nn.Parameter(torch.Tensor(1, inputs, 1, 1)) self.instance_norm_1 = nn.InstanceNorm2d(inputs, affine=False) self.blur = Blur(inputs) self.last = last self.fused_scale = fused_scale if last: self.dense = ln.Linear(inputs * 4 * 4, outputs) else: if fused_scale: self.conv_2 = ln.Conv2d(inputs, outputs, 3, 2, 1, bias=False, transform_kernel=True) else: self.conv_2 = ln.Conv2d(inputs, outputs, 3, 1, 1, bias=False) self.bias_2 = nn.Parameter(torch.Tensor(1, outputs, 1, 1)) self.instance_norm_2 = nn.InstanceNorm2d(outputs, affine=False) self.style_1 = ln.Linear(2 * inputs, latent_size) if last: self.style_2 = ln.Linear(outputs, latent_size) else: self.style_2 = ln.Linear(2 * outputs, latent_size) with torch.no_grad(): self.bias_1.zero_() self.bias_2.zero_()
def __init__(self, mapping_layers=5, latent_size=256, dlatent_size=256, mapping_fmaps=256): super(VAEMappingToLatent_old, self).__init__() inputs = latent_size self.mapping_layers = mapping_layers self.map_blocks: nn.ModuleList[MappingBlock] = nn.ModuleList() for i in range(mapping_layers): outputs = 2 * dlatent_size if i == mapping_layers - 1 else mapping_fmaps block = ln.Linear(inputs, outputs, lrmul=0.1) inputs = outputs self.map_blocks.append(block)
def __init__(self, inputs, outputs, last=False, fused_scale=True, dense=False): super(DiscriminatorBlock, self).__init__() self.conv_1 = ln.Conv2d(inputs + (1 if last else 0), inputs, 3, 1, 1, bias=False) self.bias_1 = nn.Parameter(torch.Tensor(1, inputs, 1, 1)) self.blur = Blur(inputs) self.last = last self.dense_ = dense self.fused_scale = fused_scale if self.dense_: self.dense = ln.Linear(inputs * 4 * 4, outputs) else: if fused_scale: self.conv_2 = ln.Conv2d(inputs, outputs, 3, 2, 1, bias=False, transform_kernel=True) else: self.conv_2 = ln.Conv2d(inputs, outputs, 3, 1, 1, bias=False) self.bias_2 = nn.Parameter(torch.Tensor(1, outputs, 1, 1)) with torch.no_grad(): self.bias_1.zero_() self.bias_2.zero_()
def __init__(self, startf, maxf, layer_count, latent_size, channels=3): super(EncoderWithFC, self).__init__() self.maxf = maxf self.startf = startf self.layer_count = layer_count self.from_rgb: nn.ModuleList[FromRGB] = nn.ModuleList() self.channels = channels self.latent_size = latent_size mul = 2 inputs = startf self.encode_block: nn.ModuleList[EncodeBlock] = nn.ModuleList() resolution = 2**(self.layer_count + 1) for i in range(self.layer_count): outputs = min(self.maxf, startf * mul) self.from_rgb.append(FromRGB(channels, inputs)) fused_scale = resolution >= 128 block = EncodeBlock(inputs, outputs, latent_size, i == self.layer_count - 1, fused_scale=fused_scale) resolution //= 2 #print("encode_block%d %s styles out: %d" % ((i + 1), millify(count_parameters(block)), inputs)) self.encode_block.append(block) inputs = outputs mul *= 2 self.fc2 = ln.Linear(inputs, 1, gain=1)
def __init__(self, inputs, output, lrmul): super(MappingBlock, self).__init__() self.fc = ln.Linear(inputs, output, lrmul=lrmul)