def __init__(self, config, weight_init):
super().__init__()
with tf.device('{}:*'.format(config.device)):
with tf.name_scope('discriminator_f'):
self.weight_init = weight_init
self.ch = config.gen_disc_ch
if config.conditional:
self.emb = tf.keras.layers.Embedding(
config.num_classes, self.ch * 4)
self.res_down0 = ops.resblock_down(
channels=self.ch,
config=config,
weight_init=self.weight_init,
use_bias=False)
self.bn0 = BatchNormalization()
self.att0 = ops.Attention(ch=self.ch, config=config)
self.res_down1 = ops.resblock_down(
channels=self.ch * 2,
config=config,
weight_init=self.weight_init,
use_bias=False)
self.res_down2 = ops.resblock_down(
channels=self.ch * 4,
config=config,
weight_init=self.weight_init,
use_bias=False)
self.res0 = ops.resblock(channels=self.ch * 4,
config=config,
weight_init=self.weight_init,
use_bias=False)
self.dense0 = Dense(units=1,
use_bias=True,
kernel_initializer=self.weight_init)
def __init__(self, config, weight_init):
super().__init__()
with tf.device('{}:*'.format(config.device)):
with tf.name_scope('generator'):
if config.dataset in ['mnist', 'fashion_mnist'
]: # color channel of the dataset
self.c_dim = 1
else:
self.c_dim = 3
self.weight_init = weight_init
self.num_cont_noise = config.num_cont_noise
self.ch = config.gen_disc_ch * 4
if config.conditional:
self.emb = tf.keras.layers.Embedding(
config.num_classes, config.num_emb)
self.dense0 = ops.SpectralNormalization(Dense(
units=4 * 4 * self.ch,
use_bias=True,
kernel_initializer=self.weight_init,
kernel_regularizer=ops.weight_regularizer_fully),
config=config)
self.res_up0 = ops.resblock_up_condition_top(
channels=self.ch,
config=config,
weight_init=self.weight_init,
use_bias=False)
self.res_up1 = ops.resblock_up_condition(
channels=self.ch // 2,
config=config,
weight_init=self.weight_init,
use_bias=False)
self.bn0 = BatchNormalization()
self.att0 = ops.Attention(ch=self.ch // 2, config=config)
self.res_up2 = ops.resblock_up_condition(
channels=self.ch // 4,
config=config,
weight_init=self.weight_init,
use_bias=False)
self.bn1 = BatchNormalization()
self.conv0 = ops.SpectralNormalization(Conv2DTranspose(
filters=self.c_dim,
kernel_size=3,
strides=1,
padding='SAME',
use_bias=False,
kernel_initializer=self.weight_init,
kernel_regularizer=ops.weight_regularizer),
config=config)