| 42 | |
| 43 | |
| 44 | class G(tl.Module): |
| 45 | |
| 46 | def call(self, |
| 47 | zs, |
| 48 | eps, |
| 49 | dim_10=4, |
| 50 | n_channels=3, |
| 51 | weight_norm='none', |
| 52 | feature_norm='none', |
| 53 | act=tf.nn.leaky_relu, |
| 54 | use_gram_schmidt=True, |
| 55 | training=True): |
| 56 | MAX_DIM = 512 |
| 57 | nd = lambda size: min(int(2**(10 - np.log2(size)) * dim_10), MAX_DIM) |
| 58 | |
| 59 | w_norm = tl.get_weight_norm(weight_norm, training) |
| 60 | transposed_w_norm = tl.get_weight_norm(weight_norm, training, transposed=True) |
| 61 | fc = functools.partial(tl.fc, weights_initializer=tl.get_initializer(act), weights_normalizer_fn=w_norm, weights_regularizer=slim.l2_regularizer(1.0)) |
| 62 | conv = functools.partial(tl.conv2d, weights_initializer=tl.get_initializer(act), weights_normalizer_fn=w_norm, weights_regularizer=slim.l2_regularizer(1.0)) |
| 63 | dconv = functools.partial(tl.dconv2d, weights_initializer=tl.get_initializer(act), weights_normalizer_fn=transposed_w_norm, weights_regularizer=slim.l2_regularizer(1.0)) |
| 64 | f_norm = tl.get_feature_norm(feature_norm, training, updates_collections=None) |
| 65 | f_norm = (lambda x: x) if f_norm is None else f_norm |
| 66 | |
| 67 | def orthogonal_regularizer(U): |
| 68 | with tf.name_scope('orthogonal_regularizer'): |
| 69 | U = tf.reshape(U, [-1, U.shape[-1]]) |
| 70 | orth = tf.matmul(tf.transpose(U), U) |
| 71 | tf.add_to_collections(['orth'], orth) |
| 72 | return 0.5 * tf.reduce_sum((orth - tf.eye(U.shape[-1].value)) ** 2) |
| 73 | |
| 74 | h = fc(eps, 4 * 4 * nd(4)) |
| 75 | h = tf.reshape(h, [-1, 4, 4, nd(4)]) |
| 76 | |
| 77 | for i, z in enumerate(zs): |
| 78 | height = width = 4 * 2 ** i |
| 79 | |
| 80 | U = tf.get_variable('U_%d' % i, |
| 81 | shape=[height, width, nd(height), z.shape[-1]], |
| 82 | initializer=tf.initializers.orthogonal(), |
| 83 | regularizer=orthogonal_regularizer, |
| 84 | trainable=True) |
| 85 | if use_gram_schmidt: |
| 86 | U = tf.transpose(tf.reshape(U, [-1, U.shape[-1]])) |
| 87 | U = tl.gram_schmidt(U) |
| 88 | U = tf.reshape(tf.transpose(U), [height, width, nd(height), z.shape[-1]]) |
| 89 | |
| 90 | L = tf.get_variable('L_%d' % i, |
| 91 | shape=[z.shape[-1]], |
| 92 | initializer=tf.initializers.constant([3 * i for i in range(z.shape[-1], 0, -1)]), |
| 93 | trainable=True) |
| 94 | |
| 95 | mu = tf.get_variable('mu_%d' % i, |
| 96 | shape=[height, width, nd(height)], |
| 97 | initializer=tf.initializers.zeros(), |
| 98 | trainable=True) |
| 99 | |
| 100 | h_ = tf.reduce_sum(U[None, ...] * (L[None, :] * z)[:, None, None, None, :], axis=-1) + mu[None, ...] |
| 101 | |