def define_graph(self): """ Sets up the model graph in TensorFlow. """ with tf.name_scope('discriminator'): ## # Setup scale networks. Each will make the predictions for images at a given scale. ## self.scale_nets = [] #for scale_num in xrange(self.num_scale_nets): for scale_num in range(self.num_scale_nets): with tf.name_scope('scale_net_' + str(scale_num)): scale_factor = 1. / 2**( (self.num_scale_nets - 1) - scale_num) self.scale_nets.append( DScaleModel(scale_num, int(self.height * scale_factor), int(self.width * scale_factor), self.scale_conv_layer_fms[scale_num], self.scale_kernel_sizes[scale_num], self.scale_fc_layer_sizes[scale_num])) # A list of the prediction tensors for each scale network self.scale_preds = [] #for scale_num in xrange(self.num_scale_nets): for scale_num in range(self.num_scale_nets): self.scale_preds.append(self.scale_nets[scale_num].preds) ## # Data ## self.labels = tf.placeholder(tf.float32, shape=[None, 1], name='labels') ## # Training ## with tf.name_scope('training'): # global loss is the combined loss from every scale network self.global_loss = adv_loss(self.scale_preds, self.labels) self.global_step = tf.Variable(0, trainable=False, name='global_step') self.optimizer = tf.train.GradientDescentOptimizer( c.LRATE_D, name='optimizer') self.train_op = self.optimizer.minimize( self.global_loss, global_step=self.global_step, name='train_op') # add summaries to visualize in TensorBoard #loss_summary = tf.scalar_summary('loss_D', self.global_loss) loss_summary = tf.summary.scalar('loss_D', self.global_loss) self.summaries = tf.summary.merge([loss_summary])
def define_graph(self, generator): """ Sets up the model graph in TensorFlow. @param generator: The generator model that generates frames for this to discriminate. """ with tf.name_scope('discriminator'): ## # Data ## self.input_clips = tf.placeholder( tf.float32, shape=[None, self.height, self.width, (c.HIST_LEN + c.GT_LEN) * 3]) self.g_input_frames = self.input_clips[:, :, :, :c.HIST_LEN * 3] self.gt_frames = self.input_clips[:, :, :, c.HIST_LEN * 3:] input_shape = tf.shape(self.g_input_frames) batch_size = input_shape[0] ## # Get Generator Frames ## with tf.name_scope('gen_frames'): self.g_scale_preds = [] # the generated images at each scale self.scale_gts = [] # the ground truth images at each scale self.resized_inputs = [] # the resized input images at each scale for scale_num in xrange(self.num_scale_nets): with tf.name_scope('scale_' + str(scale_num)): # for all scales but the first, add the frame generated by the last # scale to the input if scale_num > 0: last_scale_pred = self.g_scale_preds[scale_num - 1] else: last_scale_pred = None # calculate train_preds, train_gts = generator.generate_predictions(scale_num, self.height, self.width, self.g_input_frames, self.gt_frames, last_scale_pred, 'test') input_scale_factor = 1. / self.inverse_scale_factor[scale_num] input_scale_height = int(self.height * input_scale_factor) input_scale_width = int(self.width * input_scale_factor) resized_inputs = tf.image.resize_images(self.g_input_frames, [input_scale_height, input_scale_width]) self.g_scale_preds.append(train_preds) self.scale_gts.append(train_gts) self.resized_inputs.append(resized_inputs) # concatenate the generated images and ground truths at each scale self.scale_inputs = [] for scale_num in xrange(self.num_scale_nets): self.scale_inputs.append( tf.concat([self.g_scale_preds[scale_num], self.scale_gts[scale_num]], 0)) # create the labels self.labels = tf.concat([tf.zeros([batch_size, 1]), tf.ones([batch_size, 1])], 0) ## # Calculation ## # A list of the prediction tensors for each scale network self.scale_preds = [] for scale_num in xrange(self.num_scale_nets): with tf.name_scope('scale_' + str(scale_num)): with tf.name_scope('calculation'): # get predictions from the scale network self.scale_preds.append( self.scale_nets[scale_num].generate_all_predictions( tf.concat([self.resized_inputs[scale_num], self.resized_inputs[scale_num]], 0), self.scale_inputs[scale_num])) ## # Training ## with tf.name_scope('training'): # global loss is the combined loss from every scale network self.global_loss = adv_loss(self.scale_preds, self.labels) with tf.name_scope('train_step'): self.global_step = tf.Variable(0, trainable=False, name='global_step') self.optimizer = tf.train.GradientDescentOptimizer(self.lrate, name='optimizer') self.train_op = self.optimizer.minimize(self.global_loss, var_list=self.train_vars, name='train_op', global_step=self.global_step) # Accuracy test all_preds = tf.stack(self.scale_preds) self.accuracy = tf.reduce_mean(tf.cast(tf.equal(tf.round(all_preds), self.labels), tf.int32)) # add summaries to visualize in TensorBoard loss_summary = tf.summary.scalar('loss_D', self.global_loss) accuracy_summary = tf.summary.scalar('accuracy_D', self.accuracy) self.summaries = tf.summary.merge([loss_summary, accuracy_summary])
def define_graph(self): """ Sets up the model graph in TensorFlow. """ with tf.name_scope('generator'): ## # Data ## with tf.name_scope('data'): self.input_frames_train = tf.placeholder(tf.float32, shape=[ None, self.height_train, self.width_train, 3 * c.HIST_LEN ]) self.gt_frames_train = tf.placeholder( tf.float32, shape=[None, self.height_train, self.width_train, 3]) self.input_frames_test = tf.placeholder(tf.float32, shape=[ None, self.height_test, self.width_test, 3 * c.HIST_LEN ]) self.gt_frames_test = tf.placeholder( tf.float32, shape=[None, self.height_test, self.width_test, 3]) # use variable batch_size for more flexibility self.batch_size_train = tf.shape(self.input_frames_train)[0] self.batch_size_test = tf.shape(self.input_frames_test)[0] ## # Scale network setup and calculation ## self.summaries_train = [] self.scale_preds_train = [] # the generated images at each scale self.scale_gts_train = [] # the ground truth images at each scale self.d_scale_preds = [ ] # the predictions from the discriminator model self.summaries_test = [] self.scale_preds_test = [] # the generated images at each scale self.scale_gts_test = [] # the ground truth images at each scale self.scale_lastinputs_train = [] for scale_num in xrange(self.num_scale_nets): with tf.name_scope('scale_' + str(scale_num)): with tf.name_scope('setup'): ws = [] bs = [] # create weights for kernels for i in xrange(len( self.scale_kernel_sizes[scale_num])): ws.append( w([ self.scale_kernel_sizes[scale_num][i], self.scale_kernel_sizes[scale_num][i], self.scale_layer_fms[scale_num][i], self.scale_layer_fms[scale_num][i + 1] ])) bs.append( b([self.scale_layer_fms[scale_num][i + 1]])) with tf.name_scope('calculation'): def calculate(height, width, inputs, gts, last_gen_frames): # scale inputs and gts scale_factor = 1. / 2**( (self.num_scale_nets - 1) - scale_num) scale_height = int(height * scale_factor) scale_width = int(width * scale_factor) inputs = tf.image.resize_images( inputs, [scale_height, scale_width]) scale_last_inputs = inputs[:, :, :, -3:] scale_gts = tf.image.resize_images( gts, [scale_height, scale_width]) # for all scales but the first, add the frame generated by the last # scale to the input if scale_num > 0: last_gen_frames = tf.image.resize_images( last_gen_frames, [scale_height, scale_width]) inputs = tf.concat([inputs, last_gen_frames], 3) # generated frame predictions preds = inputs # perform convolutions with tf.name_scope('convolutions'): for i in xrange( len(self.scale_kernel_sizes[scale_num]) ): # Convolve layer preds = tf.nn.conv2d(preds, ws[i], [1, 1, 1, 1], padding=c.PADDING_G) # Activate with ReLU (or Tanh for last layer) if i == len( self.scale_kernel_sizes[scale_num] ) - 1: preds = tf.nn.tanh(preds + bs[i]) else: preds = tf.nn.relu(preds + bs[i]) return preds, scale_gts, scale_last_inputs ## # Perform train calculation ## # for all scales but the first, add the frame generated by the last # scale to the input if scale_num > 0: last_scale_pred_train = self.scale_preds_train[ scale_num - 1] else: last_scale_pred_train = None # calculate train_preds, train_gts, train_last_inputs = calculate( self.height_train, self.width_train, self.input_frames_train, self.gt_frames_train, last_scale_pred_train) self.scale_preds_train.append(train_preds) #self.buffer.append(train_preds) self.scale_gts_train.append(train_gts) self.scale_lastinputs_train.append(train_last_inputs) # We need to run the network first to get generated frames, run the # discriminator on those frames to get d_scale_preds, then run this # again for the loss optimization. if c.ADVERSARIAL: self.d_scale_preds.append( tf.placeholder(tf.float32, [None, 1])) ## # Perform test calculation ## # for all scales but the first, add the frame generated by the last # scale to the input if scale_num > 0: last_scale_pred_test = self.scale_preds_test[ scale_num - 1] else: last_scale_pred_test = None # calculate test_preds, test_gts, test_last_inputs = calculate( self.height_test, self.width_test, self.input_frames_test, self.gt_frames_test, last_scale_pred_test) self.scale_preds_test.append(test_preds) self.scale_gts_test.append(test_gts) ## # Training ## with tf.name_scope('train'): # global loss is the combined loss from every scale network self.global_loss = combined_loss( self.scale_preds_train, self.scale_gts_train, self.d_scale_preds, self.scale_lastinputs_train, c.L_NUM, c.ALPHA_NUM, c.LAM_ADV, c.LAM_LP, c.LAM_GDL, c.LAM_TV) self.l1_loss = c.LAM_LP * lp_loss(self.scale_preds_train, self.scale_gts_train, 1) batch_size = tf.shape(self.scale_preds_train[0])[0] self.adv_loss = c.LAM_ADV * adv_loss( self.d_scale_preds, tf.ones([4 * batch_size, 1])) self.gdl_loss = c.LAM_GDL * gdl_loss(self.scale_preds_train, self.scale_gts_train, 1) self.tv_loss = c.LAM_TV * tv_loss(self.scale_preds_train) self.global_step = tf.Variable(0, trainable=False) self.optimizer = tf.train.AdamOptimizer( learning_rate=c.LRATE_G, name='optimizer') self.train_op = self.optimizer.minimize( self.global_loss, global_step=self.global_step, name='train_op') # train loss summary loss_summary = tf.summary.scalar('train_loss_G', self.global_loss) self.summaries_train.append(loss_summary)
def train_step(self, batch, loss_list, discriminator=None): """ Runs a training step using the global loss on each of the scale networks. @param batch: An array of shape [c.BATCH_SIZE x self.height x self.width x (3 * (c.HIST_LEN + 1))]. The input and output frames, concatenated along the channel axis (index 3). @param discriminator: The discriminator model. Default = None, if not adversarial. @return: The global step. """ ## # Split into inputs and outputs ## input_frames = batch[:, :, :, :-3] gt_frames = batch[:, :, :, -3:] ## # Train ## feed_dict = { self.input_frames_train: input_frames, self.gt_frames_train: gt_frames } if c.ADVERSARIAL: # Run the generator first to get generated frames scale_preds = self.sess.run(self.scale_preds_train, feed_dict=feed_dict) # Run the discriminator nets on those frames to get predictions d_feed_dict = {} for scale_num, gen_frames in enumerate(scale_preds): d_feed_dict[discriminator.scale_nets[scale_num]. input_frames] = gen_frames d_scale_preds = self.sess.run(discriminator.scale_preds, feed_dict=d_feed_dict) d_scale_preds_list = [] # Add discriminator predictions to the for i, preds in enumerate(d_scale_preds): d_scale_preds_list.append(preds) # Calculate current loss from adversary batch_size = len(d_scale_preds_list[0]) discriminator_curr_loss = adv_loss(d_scale_preds_list, tf.ones([batch_size, 1])) loss_value = self.sess.run(discriminator_curr_loss) print(loss_value) loss_list.append(loss_value) feed_dict[self.adv_windowed_loss] = sum(loss_list) / len(loss_list) _, global_loss, global_psnr_error, global_sharpdiff_error, global_step, summaries = \ self.sess.run([self.train_op, self.global_loss, self.psnr_error_train, self.sharpdiff_error_train, self.global_step, self.summaries_train], feed_dict=feed_dict) ## # User output ## if global_step % c.STATS_FREQ == 0: print 'GeneratorModel : Step ', global_step print ' Global Loss : ', global_loss print ' PSNR Error : ', global_psnr_error print ' Sharpdiff Error: ', global_sharpdiff_error if global_step % c.SUMMARY_FREQ == 0: self.summary_writer.add_summary(summaries, global_step) print 'GeneratorModel: saved summaries' if global_step % c.IMG_SAVE_FREQ == 0: print '-' * 30 print 'Saving images...' # if not adversarial, we didn't get the preds for each scale net before for the # discriminator prediction, so do it now if not c.ADVERSARIAL: scale_preds = self.sess.run(self.scale_preds_train, feed_dict=feed_dict) # re-generate scale gt_frames to avoid having to run through TensorFlow. scale_gts = [] for scale_num in xrange(self.num_scale_nets): scale_factor = 1. / 2**((self.num_scale_nets - 1) - scale_num) scale_height = int(self.height_train * scale_factor) scale_width = int(self.width_train * scale_factor) # resize gt_output_frames for scale and append to scale_gts_train scaled_gt_frames = np.empty( [c.BATCH_SIZE, scale_height, scale_width, 3]) for i, img in enumerate(gt_frames): # for skimage.transform.resize, images need to be in range [0, 1], so normalize # to [0, 1] before resize and back to [-1, 1] after sknorm_img = (img / 2) + 0.5 resized_frame = resize(sknorm_img, [scale_height, scale_width, 3]) scaled_gt_frames[i] = (resized_frame - 0.5) * 2 scale_gts.append(scaled_gt_frames) # for every clip in the batch, save the inputs, scale preds and scale gts for pred_num in xrange(len(input_frames)): pred_dir = c.get_dir( os.path.join(c.IMG_SAVE_DIR, 'Step_' + str(global_step), str(pred_num))) # save input images for frame_num in xrange(c.HIST_LEN): img = input_frames[pred_num, :, :, (frame_num * 3):((frame_num + 1) * 3)] imsave( os.path.join(pred_dir, 'input_' + str(frame_num) + '.png'), img) # save preds and gts at each scale # noinspection PyUnboundLocalVariable for scale_num, scale_pred in enumerate(scale_preds): gen_img = scale_pred[pred_num] path = os.path.join(pred_dir, 'scale' + str(scale_num)) gt_img = scale_gts[scale_num][pred_num] imsave(path + '_gen.png', gen_img) imsave(path + '_gt.png', gt_img) print 'Saved images!' print '-' * 30 return global_step, loss_list
def define_graph(self): """ Sets up the model graph in TensorFlow. """ with tf.name_scope('discriminator'): ## # Setup scale networks. Each will make the predictions for images at a given scale. ## self.scale_nets = [] for scale_num in range(self.num_scale_nets): with tf.name_scope('scale_net_' + str(scale_num)): scale_factor = 1. / 2**( (self.num_scale_nets - 1) - scale_num) self.scale_nets.append( DScaleModel(scale_num, int(self.height * scale_factor), int(self.width * scale_factor), self.scale_conv_layer_fms[scale_num], self.scale_kernel_sizes[scale_num], self.scale_fc_layer_sizes[scale_num], self.is_w)) # A list of the prediction tensors for each scale network self.scale_preds = [] for scale_num in range(self.num_scale_nets): self.scale_preds.append(self.scale_nets[scale_num].preds) ## # Data ## self.labels = tf.placeholder(tf.float32, shape=[None, 1], name='labels') ## # Training ## with tf.name_scope('training'): # global loss is the combined loss from every scale network self.global_loss = adv_loss(self.scale_preds, self.labels, self.is_w) self.global_step = tf.Variable(0.0, trainable=False, name='global_step') with tf.control_dependencies( tf.get_collection(tf.GraphKeys.UPDATE_OPS)): if self.is_w == True: self.optimizer = tf.train.RMSPropOptimizer( self.c.LRATE_D, name='optimizer') else: self.optimizer = tf.train.GradientDescentOptimizer( self.c.LRATE_D, name='optimizer') # self.optimizer = tf.train.GradientDescentOptimizer(self.c.LRATE_D, name='optimizer') self.train_op = self.optimizer.minimize( self.global_loss, global_step=self.global_step, name='train_op', var_list=self.discriminator_vars) # add summaries to visualize in TensorBoard self.clip = [ v.assign(tf.clip_by_value(v, -0.01, 0.01)) for v in self.discriminator_vars ] loss_summary = tf.summary.scalar('loss_D', self.global_loss) self.summaries = tf.summary.merge([loss_summary])
def define_graph(self): """ Sets up the model graph in TensorFlow. """ with tf.name_scope('discriminator'): ## # Setup scale networks. Each will make the predictions for images at a given scale. ## self.scale_nets = [] for scale_num in range(self.num_scale_nets): with tf.name_scope('scale_net_' + str(scale_num)): scale_factor = 1. / 2 ** ((self.num_scale_nets - 1) - scale_num) self.scale_nets.append(DScaleModel(scale_num, int(self.height * scale_factor), int(self.width * scale_factor), self.scale_conv_layer_fms[scale_num], self.scale_kernel_sizes[scale_num], self.scale_fc_layer_sizes[scale_num])) # A list of the prediction tensors for each scale network self.scale_preds = [] for scale_num in range(self.num_scale_nets): self.scale_preds.append(self.scale_nets[scale_num].preds) ## # Data ## self.labels = tf.placeholder(tf.float32, shape=[None, 1], name='labels') ## # Training ## with tf.name_scope('training'): # global loss is the combined loss from every scale network self.global_loss = adv_loss(self.scale_preds, self.labels) if c.WASSERSTEIN and c.W_GP: epsilon = tf.random_uniform([], 0.0, 1.0) grad_penality = [] for scale_net in self.scale_nets: fake, real = tf.split(scale_net.input_frames,2) self.x_hat = real * epsilon + (1 - epsilon) * fake self.d_hat = scale_net.generate_predictions(self.x_hat) grad_penality.append(grad_penality_loss(self.x_hat, self.d_hat)) self.global_loss += c.LAM_GP * tf.reduce_mean(grad_penality) self.global_step = tf.Variable(0, trainable=False, name='global_step') self.optimizer = tf.train.AdamOptimizer(c.LRATE_D, name='optimizer') self.train_op_ = self.optimizer.minimize(self.global_loss, global_step=self.global_step, name='train_op') # Clipping to enforce 1-Lipschitz function if c.WASSERSTEIN and not c.W_GP: with tf.control_dependencies([self.train_op_]): self.train_op = tf.group(*(tf.assign(var, tf.clip_by_value(var, -c.W_Clip, c.W_Clip)) for var in tf.trainable_variables() if 'discriminator' in var.name)) else: self.train_op = self.train_op_ # add summaries to visualize in TensorBoard loss_summary = tf.summary.scalar('loss_D', self.global_loss) self.summaries = tf.summary.merge([loss_summary])