def __init__(self, n_f0, n_f1, reduce_height, in_h_res, in_w_res,
**kwargs):
super(DiscriminatorBlock, self).__init__(**kwargs)
self.gain = 1.0
self.lrmul = 1.0
self.n_f0 = n_f0
self.n_f1 = n_f1
self.reduce_height = reduce_height
self.in_h_res = in_h_res
self.in_w_res = in_w_res
self.resnet_scale = 1.0 / tf.sqrt(2.0)
# conv_0
self.conv_0 = Conv2D(
in_fmaps=self.n_f0,
out_fmaps=self.n_f0,
kernel=3,
down=False,
resample_kernel=None,
gain=self.gain,
lrmul=self.lrmul,
name="conv_0",
)
self.apply_bias_act_0 = BiasAct(lrmul=self.lrmul,
act="lrelu",
name="bias_0")
# conv_1 down
self.conv_1 = Conv2D(
in_fmaps=self.n_f0,
out_fmaps=self.n_f1,
kernel=3,
down=True,
resample_kernel=[1, 3, 3, 1],
gain=self.gain,
lrmul=self.lrmul,
name="conv_1",
reduce_height=self.reduce_height,
in_h_res=self.in_h_res,
in_w_res=self.in_w_res,
)
self.apply_bias_act_1 = BiasAct(lrmul=self.lrmul,
act="lrelu",
name="bias_1")
# resnet skip
self.conv_skip = Conv2D(
in_fmaps=self.n_f0,
out_fmaps=self.n_f1,
kernel=1,
down=True,
resample_kernel=[1, 3, 3, 1],
gain=self.gain,
lrmul=self.lrmul,
name="skip",
reduce_height=self.reduce_height,
in_h_res=self.in_h_res,
in_w_res=self.in_w_res,
)
def __init__(self, n_f0, n_f1, **kwargs):
super(DiscriminatorLastBlock, self).__init__(**kwargs)
self.gain = 1.0
self.lrmul = 1.0
self.n_f0 = n_f0
self.n_f1 = n_f1
self.minibatch_std = MinibatchStd(group_size=4,
num_new_features=1,
name="minibatchstd")
# conv_0
self.conv_0 = Conv2D(
in_fmaps=self.n_f0 + 1,
out_fmaps=self.n_f0,
kernel=3,
down=False,
resample_kernel=None,
gain=self.gain,
lrmul=self.lrmul,
name="conv_0",
)
self.apply_bias_act_0 = BiasAct(lrmul=self.lrmul,
act="lrelu",
name="bias_0")
# dense_1
self.dense_1 = Dense(self.n_f1,
gain=self.gain,
lrmul=self.lrmul,
name="dense_1")
self.apply_bias_act_1 = BiasAct(lrmul=self.lrmul,
act="lrelu",
name="bias_1")
def __init__(self, in_ch, out_fmaps, out_h_res, out_w_res, kernel_shape,
**kwargs):
super(SynthesisBlock, self).__init__(**kwargs)
self.in_ch = in_ch
self.fmaps = out_fmaps
self.gain = 1.0
self.lrmul = 1.0
self.out_h_res = out_h_res
self.out_w_res = out_w_res
self.kernel_shape = kernel_shape
# conv0 up
self.conv_0 = ModulatedConv2D(
in_fmaps=self.in_ch,
out_fmaps=self.fmaps,
kernel_shape=self.kernel_shape,
up=True,
demodulate=True,
resample_kernel=[1, 3, 3, 1],
gain=self.gain,
lrmul=self.lrmul,
fused_modconv=True,
in_h_res=self.out_h_res // 2,
in_w_res=self.out_w_res // 2,
name="conv_0",
)
self.apply_noise_0 = Noise(name="noise_0")
self.apply_bias_act_0 = BiasAct(lrmul=self.lrmul,
act="lrelu",
name="bias_0")
# conv block
self.conv_1 = ModulatedConv2D(
in_w_res=self.out_w_res,
in_h_res=self.out_h_res,
in_fmaps=self.fmaps,
out_fmaps=self.fmaps,
kernel_shape=self.kernel_shape,
up=False,
demodulate=True,
resample_kernel=[1, 3, 3, 1],
gain=self.gain,
lrmul=self.lrmul,
fused_modconv=True,
name="conv_1",
)
self.apply_noise_1 = Noise(name="noise_1")
self.apply_bias_act_1 = BiasAct(lrmul=self.lrmul,
act="lrelu",
name="bias_1")
def __init__(self, fmaps, h_res, w_res, **kwargs):
super(FromRGB, self).__init__(**kwargs)
self.fmaps = fmaps
self.h_res = h_res
self.w_res = w_res
self.conv = Conv2D(
in_fmaps=3,
out_fmaps=self.fmaps,
kernel=1,
down=False,
resample_kernel=None,
gain=1.0,
lrmul=1.0,
name="conv",
)
self.apply_bias_act = BiasAct(lrmul=1.0, act="lrelu", name="bias")
def __init__(self, **kwargs):
super(Discriminator, self).__init__(**kwargs)
# discriminator's (resolutions and featuremaps) are reversed against generator's
self.resolutions = cfg.discrim_resolutions
self.feat_maps = cfg.discrim_feat_maps
# stack discriminator blocks
res0, n_f0 = self.resolutions[0], self.feat_maps[0]
self.initial_fromrgb = FromRGB(
fmaps=n_f0,
h_res=res0[0],
w_res=res0[1],
name="{:d}x{:d}/FromRGB".format(res0[0], res0[1]),
)
self.blocks = list()
for res, next_step_res, f_m0, f_m1 in zip(
self.resolutions[:-1],
self.resolutions[1:],
self.feat_maps[:-1],
self.feat_maps[1:],
):
self.blocks.append(
DiscriminatorBlock(
n_f0=f_m0,
n_f1=f_m1,
reduce_height=res[0] != next_step_res[0],
in_h_res=res[0],
in_w_res=res[1],
name="{:d}x{:d}".format(res[0], res0[1]),
), )
# set last discriminator block
res_final = self.resolutions[-1]
n_f0, n_f1 = self.feat_maps[-2], self.feat_maps[-1]
self.last_block = DiscriminatorLastBlock(n_f0,
n_f1,
name="{:d}x{:d}".format(
res_final[0],
res_final[1]))
# set last dense layer
self.last_dense = Dense(1, gain=1.0, lrmul=1.0, name="last_dense")
self.last_bias = BiasAct(lrmul=1.0, act="linear", name="last_bias")
def __init__(self, in_ch, h_res, w_res, **kwargs):
super(ToRGB, self).__init__(**kwargs)
self.in_ch = in_ch
self.h_res = h_res
self.w_res = w_res
self.conv = ModulatedConv2D(
in_fmaps=self.in_ch,
out_fmaps=3,
kernel_shape=[1, 1],
up=False,
demodulate=False,
resample_kernel=None,
gain=1.0,
lrmul=1.0,
fused_modconv=True,
name="conv",
)
self.apply_bias = BiasAct(lrmul=1.0, act="linear", name="bias")
def __init__(self,
in_fmaps,
out_fmaps,
kernel_shape,
up,
demodulate,
resample_kernel,
gain,
lrmul,
fused_modconv,
in_h_res=None,
in_w_res=None,
**kwargs):
super(ModulatedConv2D, self).__init__(**kwargs)
self.in_fmaps = in_fmaps
self.out_fmaps = out_fmaps
self.kernel_shape = kernel_shape
self.demodulate = demodulate
self.up = up
self.fused_modconv = fused_modconv
self.in_h_res = in_h_res
self.in_w_res = in_w_res
self.gain = gain
self.lrmul = lrmul
self.k, self.pad0, self.pad1 = compute_paddings(resample_kernel,
up,
False,
is_conv=True)
# self.factor = 2
self.mod_dense = Dense(self.in_fmaps,
gain=1.0,
lrmul=1.0,
name="mod_dense")
self.mod_bias = BiasAct(lrmul=1.0, act="linear", name="mod_bias")
def __init__(self, style_dim, n_mapping, name, **kwargs):
super(Mapping, self).__init__(name=name, **kwargs)
self.style_dim = style_dim
self.n_mapping = n_mapping
self.gain = 1.0
self.lrmul = 0.01
self.normalize = tf.keras.layers.Lambda(lambda x: x * tf.math.rsqrt(
tf.reduce_mean(tf.square(x), axis=1, keepdims=True) + 1e-8))
self.dense_layers = list()
self.bias_act_layers = list()
for ii in range(self.n_mapping):
self.dense_layers.append(
Dense(
self.style_dim,
gain=self.gain,
lrmul=self.lrmul,
name="dense_{:d}".format(ii),
))
self.bias_act_layers.append(
BiasAct(lrmul=self.lrmul,
act="lrelu",
name="bias_{:d}".format(ii)))