def blend_images(x, progress, resolution_schedule, num_blocks): """Blends images of different resolutions according to `progress`. When training `progress` is at a stable stage for resolution r, returns image `x` downscaled to resolution r and then upscaled to `final_resolutions`, call it x'(r). Otherwise when training `progress` is at a transition stage from resolution r to 2r, returns a linear combination of x'(r) and x'(2r). Args: x: An image `Tensor` of NHWC format with resolution `final_resolutions`. progress: A scalar float `Tensor` of training progress. resolution_schedule: An object of `ResolutionSchedule`. num_blocks: An integer of number of blocks. Returns: An image `Tensor` which is a blend of images of different resolutions. """ x_blend = [] for block_id in range(1, num_blocks + 1): alpha = _generator_alpha(block_id, progress) scale = resolution_schedule.scale_factor(block_id) x_blend.append(alpha * layers.upscale(layers.downscale(x, scale), scale)) return tf.add_n(x_blend)
def __call__(self, x, x_cond=None, params=None): del params if x_cond is not None: x = tf.concat([x, x_cond], axis=3) responses = [] for ii, D in enumerate(self.discriminators): responses.append(D(x, x_cond=None)) # x_cond is already concatenated if ii != len(self.discriminators) - 1: x = layers.downscale(x, n=2) return responses
def test_downscale_4d_images_returns_downscaled_images(self): x_np = np.array( [[[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [10, 11, 12]]], [[[0, 0, 0], [-1, -2, -3]], [[1, -2, 2], [2, 5, 3]]]], dtype=np.float32) with self.test_session(use_gpu=True) as sess: x1_np, x2_np = sess.run( [layers.downscale(tf.constant(x_np), n) for n in [1, 2]]) expected2_np = [[[[5.5, 6.5, 7.5]]], [[[0.5, 0.25, 0.5]]]] self.assertNDArrayNear(x1_np, x_np, 1.0e-5) self.assertNDArrayNear(x2_np, expected2_np, 1.0e-5)
def test_blend_images_in_transition_stage(self): x_np = np.random.normal(size=[2, 8, 8, 3]) x = tf.constant(x_np, tf.float32) x_blend = networks.blend_images( x, tf.constant(0.2), resolution_schedule=networks.ResolutionSchedule(scale_base=2, num_resolutions=2), num_blocks=2) with self.test_session(use_gpu=True) as sess: x_blend_np = sess.run(x_blend) x_blend_expected_np = 0.8 * sess.run( layers.upscale(layers.downscale(x, 2), 2)) + 0.2 * x_np self.assertNDArrayNear(x_blend_np, x_blend_expected_np, 1.0e-6)
def discriminator(x, progress, num_filters_fn, resolution_schedule, num_blocks=None, kernel_size=3, scope='progressive_gan_discriminator', reuse=None): """Discriminator network for the progressive GAN model. Args: x: A `Tensor`of NHWC format representing images of size `resolution`. progress: A scalar float `Tensor` of training progress. num_filters_fn: A function that maps `block_id` to # of filters for the block. resolution_schedule: An object of `ResolutionSchedule`. num_blocks: An integer of number of blocks. None means maximum number of blocks, i.e. `resolution.schedule.num_resolutions`. Defaults to None. kernel_size: An integer of convolution kernel size. scope: A string or variable scope. reuse: Whether to reuse `scope`. Defaults to None which means to inherit the reuse option of the parent scope. Returns: A `Tensor` of model output and a dictionary of model end points. """ if num_blocks is None: num_blocks = resolution_schedule.num_resolutions def _conv2d(scope, x, kernel_size, filters, padding='SAME'): return layers.custom_conv2d(x=x, filters=filters, kernel_size=kernel_size, padding=padding, activation=tf.nn.leaky_relu, he_initializer_slope=0.0, scope=scope) def _from_rgb(x, block_id): return _conv2d('from_rgb', x, 1, num_filters_fn(block_id)) end_points = {} with tf.variable_scope(scope, reuse=reuse): x0 = x end_points['rgb'] = x0 lods = [] for block_id in range(num_blocks, 0, -1): with tf.variable_scope(block_name(block_id)): scale = resolution_schedule.scale_factor(block_id) lod = layers.downscale(x0, scale) end_points['downscaled_rgb_{}'.format(block_id)] = lod lod = _from_rgb(lod, block_id) # alpha_i is used to replace lod_select. alpha = _discriminator_alpha(block_id, progress) end_points['alpha_{}'.format(block_id)] = alpha lods.append((lod, alpha)) lods_iter = iter(lods) x, _ = lods_iter.__next__() for block_id in range(num_blocks, 1, -1): with tf.variable_scope(block_name(block_id)): x = _conv2d('conv0', x, kernel_size, num_filters_fn(block_id)) x = _conv2d('conv1', x, kernel_size, num_filters_fn(block_id - 1)) x = layers.downscale(x, resolution_schedule.scale_base) lod, alpha = lods_iter.__next__() x = alpha * lod + (1.0 - alpha) * x with tf.variable_scope(block_name(1)): x = layers.scalar_concat(x, layers.minibatch_mean_stddev(x)) x = _conv2d('conv0', x, kernel_size, num_filters_fn(1)) x = _conv2d('conv1', x, resolution_schedule.start_resolutions, num_filters_fn(0), 'VALID') end_points['last_conv'] = x logits = layers.custom_dense(x=x, units=1, scope='logits') end_points['logits'] = logits return logits, end_points
def test_downscale_invalid_scale_throws_exception(self): with self.assertRaises(ValueError): layers.downscale(tf.constant([]), -1)
def discriminator(x, progress, num_filters_fn, resolution_schedule, num_blocks=None, kernel_size=3, scope='progressive_gan_discriminator', reuse=None): """Discriminator network for the progressive GAN model. Args: x: A `Tensor`of NHWC format representing images of size `resolution`. progress: A scalar float `Tensor` of training progress. num_filters_fn: A function that maps `block_id` to # of filters for the block. resolution_schedule: An object of `ResolutionSchedule`. num_blocks: An integer of number of blocks. None means maximum number of blocks, i.e. `resolution.schedule.num_resolutions`. Defaults to None. kernel_size: An integer of convolution kernel size. scope: A string or variable scope. reuse: Whether to reuse `scope`. Defaults to None which means to inherit the reuse option of the parent scope. Returns: A `Tensor` of model output and a dictionary of model end points. """ if num_blocks is None: num_blocks = resolution_schedule.num_resolutions def _conv2d(scope, x, kernel_size, filters, padding='SAME'): return layers.custom_conv2d( x=x, filters=filters, kernel_size=kernel_size, padding=padding, activation=tf.nn.leaky_relu, he_initializer_slope=0.0, scope=scope) def _from_rgb(x, block_id): return _conv2d('from_rgb', x, 1, num_filters_fn(block_id)) end_points = {} with tf.variable_scope(scope, reuse=reuse): x0 = x end_points['rgb'] = x0 lods = [] for block_id in range(num_blocks, 0, -1): with tf.variable_scope(block_name(block_id)): scale = resolution_schedule.scale_factor(block_id) lod = layers.downscale(x0, scale) end_points['downscaled_rgb_{}'.format(block_id)] = lod lod = _from_rgb(lod, block_id) # alpha_i is used to replace lod_select. alpha = _discriminator_alpha(block_id, progress) end_points['alpha_{}'.format(block_id)] = alpha lods.append((lod, alpha)) lods_iter = iter(lods) x, _ = lods_iter.next() for block_id in range(num_blocks, 1, -1): with tf.variable_scope(block_name(block_id)): x = _conv2d('conv0', x, kernel_size, num_filters_fn(block_id)) x = _conv2d('conv1', x, kernel_size, num_filters_fn(block_id - 1)) x = layers.downscale(x, resolution_schedule.scale_base) lod, alpha = lods_iter.next() x = alpha * lod + (1.0 - alpha) * x with tf.variable_scope(block_name(1)): x = layers.scalar_concat(x, layers.minibatch_mean_stddev(x)) x = _conv2d('conv0', x, kernel_size, num_filters_fn(1)) x = _conv2d('conv1', x, resolution_schedule.start_resolutions, num_filters_fn(0), 'VALID') end_points['last_conv'] = x logits = layers.custom_dense(x=x, units=1, scope='logits') end_points['logits'] = logits return logits, end_points