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 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