Downloads: 1
Diff of
/trainers/fAnoGAN.py
[000000]
..
[978658]
Switch to unified view
| a | b/trainers/fAnoGAN.py | ||
|---|---|---|---|
| 1 | from collections import defaultdict |
||
| 2 | from math import inf |
||
| 3 | |||
| 4 | from tensorflow.python.ops.losses.losses_impl import Reduction |
||
| 5 | |||
| 6 | from trainers import trainer_utils |
||
| 7 | from trainers.AEMODEL import AEMODEL, Phase, indicate_early_stopping, update_log_dicts |
||
| 8 | from trainers.DLMODEL import * |
||
| 9 | |||
| 10 | |||
| 11 | class fAnoGAN(AEMODEL): |
||
| 12 | class Config(AEMODEL.Config): |
||
| 13 | def __init__(self): |
||
| 14 | super().__init__('fAnoGAN') |
||
| 15 | self.scale = 10.0 |
||
| 16 | self.kappa = 1.0 |
||
| 17 | |||
| 18 | def __init__(self, sess, config, network=None): |
||
| 19 | super().__init__(sess, config, network) |
||
| 20 | self.x = tf.placeholder(tf.float32, [None, self.config.outputHeight, self.config.outputWidth, self.config.numChannels], name='x') |
||
| 21 | self.z = tf.placeholder(tf.float32, [None, self.config.zDim], name='z') |
||
| 22 | |||
| 23 | self.outputs = self.network(self.z, self.x, dropout_rate=self.dropout_rate, dropout=self.dropout, config=self.config) |
||
| 24 | self.z_enc = self.outputs['z_enc'] |
||
| 25 | self.generated = self.x_ = self.outputs['x_'] |
||
| 26 | self.reconstruction = self.x_enc = self.outputs['x_enc'] |
||
| 27 | self.d_fake_features = self.outputs['d_fake_features'] |
||
| 28 | self.d_ = self.outputs['d_'] |
||
| 29 | self.d_features = self.outputs['d_features'] |
||
| 30 | self.d = self.outputs['d'] |
||
| 31 | self.x_hat = self.outputs['x_hat'] |
||
| 32 | self.d_hat_features = self.outputs['d_hat_features'] |
||
| 33 | self.d_hat = self.outputs['d_hat'] |
||
| 34 | self.d_enc_features = self.outputs['d_enc_features'] |
||
| 35 | self.d_enc = self.outputs['d_enc'] |
||
| 36 | |||
| 37 | self.kappa = self.config.kappa |
||
| 38 | self.scale = self.config.scale |
||
| 39 | |||
| 40 | # Print Stats |
||
| 41 | self.get_number_of_trainable_params() |
||
| 42 | # Instantiate Saver |
||
| 43 | self.saver = tf.train.Saver() |
||
| 44 | |||
| 45 | def train(self, dataset): |
||
| 46 | # Determine trainable variables |
||
| 47 | self.variables = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES) |
||
| 48 | |||
| 49 | # Build losses |
||
| 50 | self.losses['disc_real'] = disc_real = tf.reduce_mean(self.d) |
||
| 51 | self.losses['disc_fake'] = disc_fake = tf.reduce_mean(self.d_) |
||
| 52 | self.losses['gen_loss'] = gen_loss = -disc_fake |
||
| 53 | disc_loss = disc_fake - disc_real |
||
| 54 | |||
| 55 | ddx = tf.gradients(self.d_hat, self.x_hat)[0] # gradient |
||
| 56 | ddx = tf.sqrt(tf.reduce_sum(tf.square(ddx), axis=1)) # slopes |
||
| 57 | ddx = tf.reduce_mean(tf.square(ddx - 1.0)) * self.scale # gradient penalty |
||
| 58 | self.losses['disc_loss'] = disc_loss = disc_loss + ddx |
||
| 59 | |||
| 60 | self.losses['loss_img'] = loss_img = tf.reduce_mean( |
||
| 61 | tf.reduce_mean(tf.losses.mean_squared_error(self.x, self.x_enc, reduction=Reduction.NONE), axis=[1, 2, 3])) |
||
| 62 | self.losses['loss_fts'] = loss_fts = tf.reduce_mean( |
||
| 63 | tf.reduce_mean(tf.losses.mean_squared_error(self.d_enc_features, self.d_features, reduction=Reduction.NONE), axis=[1, 2, 3])) |
||
| 64 | self.losses['enc_loss'] = enc_loss = loss_img + self.kappa * loss_fts |
||
| 65 | self.losses['L1'] = tf.losses.absolute_difference(self.x, self.x_enc, reduction=Reduction.NONE) |
||
| 66 | self.losses['reconstructionLoss'] = self.losses['loss'] = tf.reduce_mean(tf.reduce_sum(self.losses['L1'], axis=[1, 2, 3])) |
||
| 67 | |||
| 68 | with tf.control_dependencies(tf.get_collection(tf.GraphKeys.UPDATE_OPS)): |
||
| 69 | # Set the optimizer |
||
| 70 | t_vars = tf.trainable_variables() |
||
| 71 | dis_vars = [var for var in t_vars if 'Discriminator' in var.name] |
||
| 72 | gen_vars = [var for var in t_vars if 'Generator' in var.name] |
||
| 73 | enc_vars = [var for var in t_vars if 'Encoder' in var.name] |
||
| 74 | |||
| 75 | optim_dis = tf.train.AdamOptimizer(learning_rate=self.config.learningrate, beta1=0.5, beta2=0.9).minimize(disc_loss, var_list=dis_vars) |
||
| 76 | optim_gen = tf.train.AdamOptimizer(learning_rate=self.config.learningrate, beta1=0.5, beta2=0.9).minimize(gen_loss, var_list=gen_vars) |
||
| 77 | optim_enc = tf.train.AdamOptimizer(learning_rate=self.config.learningrate, beta1=0.5, beta2=0.9).minimize(enc_loss, var_list=enc_vars) |
||
| 78 | |||
| 79 | # initialize all variables |
||
| 80 | tf.global_variables_initializer().run(session=self.sess) |
||
| 81 | |||
| 82 | best_cost = inf |
||
| 83 | last_improvement = 0 |
||
| 84 | last_epoch = self.load_checkpoint() |
||
| 85 | |||
| 86 | # Go go go! |
||
| 87 | for epoch in range(last_epoch, self.config.numEpochs): |
||
| 88 | ################# |
||
| 89 | # TRAINING WGAN # |
||
| 90 | ################# |
||
| 91 | phase = Phase.TRAIN |
||
| 92 | scalars = defaultdict(list) |
||
| 93 | visuals = [] |
||
| 94 | d_iters = 5 |
||
| 95 | num_batches = dataset.num_batches(self.config.batchsize, set=phase.value) |
||
| 96 | for idx in range(0, num_batches): |
||
| 97 | batch, _, _ = dataset.next_batch(self.config.batchsize, set=phase.value) |
||
| 98 | |||
| 99 | # Generator optimization |
||
| 100 | fetches = { |
||
| 101 | 'generated': self.generated, |
||
| 102 | 'gen_loss': self.losses['gen_loss'], |
||
| 103 | 'optimizer_g': optim_gen, |
||
| 104 | } |
||
| 105 | |||
| 106 | feed_dict = { |
||
| 107 | self.x: batch, |
||
| 108 | self.z: self.sample_z(), |
||
| 109 | self.dropout: phase == Phase.TRAIN, |
||
| 110 | self.dropout_rate: self.config.dropout_rate |
||
| 111 | } |
||
| 112 | run = self.sess.run(fetches, feed_dict=feed_dict) |
||
| 113 | |||
| 114 | for _ in range(0, d_iters): |
||
| 115 | # Discriminator optimization |
||
| 116 | fetches = { |
||
| 117 | 'generated': self.generated, |
||
| 118 | 'disc_loss': self.losses['disc_loss'], |
||
| 119 | 'disc_fake': self.losses['disc_fake'], |
||
| 120 | 'disc_real': self.losses['disc_real'], |
||
| 121 | 'optimizer_d': optim_dis, |
||
| 122 | } |
||
| 123 | feed_dict = { |
||
| 124 | self.x: batch, |
||
| 125 | self.z: self.sample_z(), |
||
| 126 | self.dropout: phase == Phase.TRAIN, |
||
| 127 | self.dropout_rate: self.config.dropout_rate |
||
| 128 | } |
||
| 129 | run = {**run, **self.sess.run(fetches, feed_dict=feed_dict)} |
||
| 130 | |||
| 131 | # Print to console |
||
| 132 | print(f'Epoch ({phase.value} WGAN): [{epoch:2d}] [{idx:4d}/{num_batches:4d}]' |
||
| 133 | f' gen_loss: {run["gen_loss"]:.8f}, disc_loss: {run["disc_loss"]:.8f}') |
||
| 134 | update_log_dicts(*trainer_utils.get_summary_dict(batch, run, visualization_keys=['generated']), scalars, visuals) |
||
| 135 | |||
| 136 | self.log_to_tensorboard(epoch, scalars, visuals, phase, name='wgan_x') |
||
| 137 | |||
| 138 | # Increment last_epoch counter and save model |
||
| 139 | last_epoch += 1 |
||
| 140 | self.save(self.checkpointDir, last_epoch) |
||
| 141 | |||
| 142 | for epoch in range(last_epoch, 2 * self.config.numEpochs): |
||
| 143 | #################### |
||
| 144 | # TRAINING Encoder # |
||
| 145 | #################### |
||
| 146 | phase = Phase.TRAIN |
||
| 147 | scalars = defaultdict(list) |
||
| 148 | visuals = [] |
||
| 149 | num_batches = dataset.num_batches(self.config.batchsize, set=phase.value) |
||
| 150 | for idx in range(0, num_batches): |
||
| 151 | batch, _, _ = dataset.next_batch(self.config.batchsize, set=phase.value) |
||
| 152 | fetches = { |
||
| 153 | 'reconstruction': self.reconstruction, |
||
| 154 | 'optimizer_enc': optim_enc, |
||
| 155 | 'z_enc': self.z_enc, |
||
| 156 | 'z': self.z, |
||
| 157 | **self.losses |
||
| 158 | } |
||
| 159 | |||
| 160 | feed_dict = { |
||
| 161 | self.x: batch, |
||
| 162 | self.z: self.sample_z(), |
||
| 163 | self.dropout: phase == Phase.TRAIN, |
||
| 164 | self.dropout_rate: self.config.dropout_rate |
||
| 165 | } |
||
| 166 | run = self.sess.run(fetches, feed_dict=feed_dict) |
||
| 167 | |||
| 168 | # Print to console |
||
| 169 | print(f'Epoch ({phase.value} Encoder): [{epoch:2d}] [{idx:4d}/{num_batches:4d}] reconstructionLoss: {run["reconstructionLoss"]:.8f}') |
||
| 170 | update_log_dicts(*trainer_utils.get_summary_dict(batch, run), scalars, visuals) |
||
| 171 | |||
| 172 | self.log_to_tensorboard(epoch, scalars, visuals, phase) |
||
| 173 | |||
| 174 | # Increment last_epoch counter and save model |
||
| 175 | last_epoch += 1 |
||
| 176 | self.save(self.checkpointDir, last_epoch) |
||
| 177 | |||
| 178 | ###################### |
||
| 179 | # VALIDATION Encoder # |
||
| 180 | ###################### |
||
| 181 | phase = Phase.VAL |
||
| 182 | scalars = defaultdict(list) |
||
| 183 | visuals = [] |
||
| 184 | num_batches = dataset.num_batches(self.config.batchsize, set=phase.value) |
||
| 185 | for idx in range(0, num_batches): |
||
| 186 | batch, _, _ = dataset.next_batch(self.config.batchsize, set=phase.value) |
||
| 187 | |||
| 188 | fetches = { |
||
| 189 | 'reconstruction': self.reconstruction, |
||
| 190 | **self.losses |
||
| 191 | } |
||
| 192 | |||
| 193 | feed_dict = { |
||
| 194 | self.x: batch, |
||
| 195 | self.z: self.sample_z(), |
||
| 196 | self.dropout: phase == Phase.TRAIN, |
||
| 197 | self.dropout_rate: self.config.dropout_rate |
||
| 198 | } |
||
| 199 | run = self.sess.run(fetches, feed_dict=feed_dict) |
||
| 200 | |||
| 201 | # Print to console |
||
| 202 | print(f'Epoch ({phase.value}): [{epoch:2d}] [{idx:4d}/{num_batches:4d}] reconstructionLoss: {run["reconstructionLoss"]:.8f}') |
||
| 203 | update_log_dicts(*trainer_utils.get_summary_dict(batch, run), scalars, visuals) |
||
| 204 | |||
| 205 | self.log_to_tensorboard(epoch, scalars, visuals, phase) |
||
| 206 | |||
| 207 | best_cost, last_improvement, stop = indicate_early_stopping(scalars['reconstructionLoss'], best_cost, last_improvement) |
||
| 208 | if stop: |
||
| 209 | print('Early stopping was triggered due to no improvement over the last 5 epochs') |
||
| 210 | break |
||
| 211 | |||
| 212 | def get_feed_dict(self, batch, phase): |
||
| 213 | return { |
||
| 214 | self.x: batch, |
||
| 215 | self.z: self.sample_z(), |
||
| 216 | self.dropout: phase == Phase.TRAIN, |
||
| 217 | self.dropout_rate: self.config.dropout_rate |
||
| 218 | } |
||
| 219 | |||
| 220 | def reconstruct(self, x, dropout=False): |
||
| 221 | if x.ndim < 4: |
||
| 222 | x = np.expand_dims(x, 0) |
||
| 223 | |||
| 224 | fetches = { |
||
| 225 | 'reconstruction': self.reconstruction |
||
| 226 | } |
||
| 227 | |||
| 228 | feed_dict = { |
||
| 229 | self.x: x, |
||
| 230 | self.z: self.sample_z(x.shape[0]), |
||
| 231 | self.dropout: dropout, # apply only during MC sampling. |
||
| 232 | self.dropout_rate: self.config.dropout_rate |
||
| 233 | } |
||
| 234 | results = self.sess.run(fetches, feed_dict=feed_dict) |
||
| 235 | |||
| 236 | results['l1err'] = np.sum(np.abs(x - results['reconstruction'])) |
||
| 237 | results['l2err'] = np.sum(np.sqrt((x - results['reconstruction']) ** 2)) |
||
| 238 | |||
| 239 | return results |
||
| 240 | |||
| 241 | def sample_z(self, batch_size=None): |
||
| 242 | return np.random.normal(size=[batch_size if batch_size else self.config.batchsize, self.config.zDim]) |