def build(config, batch_size, is_train=False):
optimizer = train_utils.build_optimizer(config)
ema_vars = []
downsample = config.get('downsample', False)
downsample_res = config.get('downsample_res', 64)
h, w = config.resolution
if config.model.name == 'coltran_core':
if downsample:
h, w = downsample_res, downsample_res
zero = tf.zeros((batch_size, h, w, 3), dtype=tf.int32)
model = colorizer.ColTranCore(config.model)
model(zero, training=is_train)
c = 1 if is_train else 3
if config.model.name == 'color_upsampler':
if downsample:
h, w = downsample_res, downsample_res
zero_slice = tf.zeros((batch_size, h, w, c), dtype=tf.int32)
zero = tf.zeros((batch_size, h, w, 3), dtype=tf.int32)
model = upsampler.ColorUpsampler(config.model)
model(zero, inputs_slice=zero_slice, training=is_train)
elif config.model.name == 'spatial_upsampler':
zero_slice = tf.zeros((batch_size, h, w, c), dtype=tf.int32)
zero = tf.zeros((batch_size, h, w, 3), dtype=tf.int32)
model = upsampler.SpatialUpsampler(config.model)
model(zero, inputs_slice=zero_slice, training=is_train)
ema_vars = model.trainable_variables
ema = train_utils.build_ema(config, ema_vars)
return model, optimizer, ema
def test_bit_upsampler_attention_num_channels_1(self):
config = self.get_config()
bit_upsampler = upsampler.ColorUpsampler(config=config)
inputs = tf.random.uniform(shape=(8, 32, 32, 3),
minval=0,
maxval=256,
dtype=tf.int32)
inputs_slice = tf.random.uniform(shape=(8, 32, 32, 1),
minval=0,
maxval=256,
dtype=tf.int32)
grayscale = tf.image.rgb_to_grayscale(inputs)
channel_index = tf.random.uniform(shape=[
8,
],
minval=0,
maxval=3,
dtype=tf.int32)
logits = bit_upsampler(inputs=inputs,
inputs_slice=inputs_slice,
channel_index=channel_index)[0]
self.assertEqual(logits.shape, (8, 32, 32, 1, 256))
inputs = base_utils.convert_bits(inputs, n_bits_in=8, n_bits_out=3)
output = bit_upsampler.sample(gray_cond=grayscale, bit_cond=inputs)
self.assertEqual(output['bit_up_argmax'].shape, (8, 32, 32, 3))
def build_model(config):
"""Builds model."""
name = config.model.name
optimizer = train_utils.build_optimizer(config)
zero_64 = tf.zeros((1, 64, 64, 3), dtype=tf.int32)
zero_64_slice = tf.zeros((1, 64, 64, 1), dtype=tf.int32)
zero = tf.zeros((1, 256, 256, 3), dtype=tf.int32)
zero_slice = tf.zeros((1, 256, 256, 1), dtype=tf.int32)
if name == 'coltran_core':
model = colorizer.ColTranCore(config.model)
model(zero_64, training=False)
elif name == 'color_upsampler':
model = upsampler.ColorUpsampler(config.model)
model(inputs=zero_64, inputs_slice=zero_64_slice, training=False)
elif name == 'spatial_upsampler':
model = upsampler.SpatialUpsampler(config.model)
model(inputs=zero, inputs_slice=zero_slice, training=False)
ema_vars = model.trainable_variables
ema = train_utils.build_ema(config, ema_vars)
return model, optimizer, ema
