| 23 | |
| 24 | |
| 25 | class Model(DCGAN.Model): |
| 26 | # replace BatchNorm by LayerNorm |
| 27 | @auto_reuse_variable_scope |
| 28 | def discriminator(self, imgs): |
| 29 | nf = 64 |
| 30 | with argscope(Conv2D, activation=tf.identity, kernel_size=4, strides=2): |
| 31 | l = (LinearWrap(imgs) |
| 32 | .Conv2D('conv0', nf, activation=tf.nn.leaky_relu) |
| 33 | .Conv2D('conv1', nf * 2) |
| 34 | .LayerNorm('ln1') |
| 35 | .tf.nn.leaky_relu() |
| 36 | .Conv2D('conv2', nf * 4) |
| 37 | .LayerNorm('ln2') |
| 38 | .tf.nn.leaky_relu() |
| 39 | .Conv2D('conv3', nf * 8) |
| 40 | .LayerNorm('ln3') |
| 41 | .tf.nn.leaky_relu() |
| 42 | .FullyConnected('fct', 1, activation=tf.identity)()) |
| 43 | return tf.reshape(l, [-1]) |
| 44 | |
| 45 | def build_graph(self, image_pos): |
| 46 | image_pos = image_pos / 128.0 - 1 |
| 47 | |
| 48 | z = tf.random_normal([self.batch, self.zdim], name='z_train') |
| 49 | z = tf.placeholder_with_default(z, [None, self.zdim], name='z') |
| 50 | |
| 51 | with argscope([Conv2D, Conv2DTranspose, FullyConnected], |
| 52 | kernel_initializer=tf.truncated_normal_initializer(stddev=0.02)): |
| 53 | with tf.variable_scope('gen'): |
| 54 | image_gen = self.generator(z) |
| 55 | tf.summary.image('generated-samples', image_gen, max_outputs=30) |
| 56 | |
| 57 | alpha = tf.random_uniform(shape=[self.batch, 1, 1, 1], |
| 58 | minval=0., maxval=1., name='alpha') |
| 59 | interp = image_pos + alpha * (image_gen - image_pos) |
| 60 | |
| 61 | with tf.variable_scope('discrim'): |
| 62 | vecpos = self.discriminator(image_pos) |
| 63 | vecneg = self.discriminator(image_gen) |
| 64 | vec_interp = self.discriminator(interp) |
| 65 | |
| 66 | # the Wasserstein-GAN losses |
| 67 | self.d_loss = tf.reduce_mean(vecneg - vecpos, name='d_loss') |
| 68 | self.g_loss = tf.negative(tf.reduce_mean(vecneg), name='g_loss') |
| 69 | |
| 70 | # the gradient penalty loss |
| 71 | gradients = tf.gradients(vec_interp, [interp])[0] |
| 72 | gradients = tf.sqrt(tf.reduce_sum(tf.square(gradients), [1, 2, 3])) |
| 73 | gradients_rms = tf.sqrt(tf.reduce_mean(tf.square(gradients)), name='gradient_rms') |
| 74 | gradient_penalty = tf.reduce_mean(tf.square(gradients - 1), name='gradient_penalty') |
| 75 | add_moving_summary(self.d_loss, self.g_loss, gradient_penalty, gradients_rms) |
| 76 | |
| 77 | self.d_loss = tf.add(self.d_loss, 10 * gradient_penalty) |
| 78 | |
| 79 | self.collect_variables() |
| 80 | |
| 81 | def optimizer(self): |
| 82 | opt = tf.train.AdamOptimizer(1e-4, beta1=0.5, beta2=0.9) |