def define_graph(self): """ Sets up the model graph in TensorFlow. """ with tf.name_scope('generator'): ## # Data ## with tf.name_scope('data'): self.inputs = tf.placeholder(tf.float32, shape=[None, 6]) self.gt_frames = tf.placeholder( tf.float32, shape=[None, self.height, self.width, 3]) # use variable batch_size for more flexibility self.batch_size = tf.shape(self.inputs)[0] ## # Scale network setup and calculation ## self.summaries = [] self.scale_preds = [] # the generated images at each scale self.scale_gts = [] # the ground truth images at each scale self.d_scale_preds = [ ] # the predictions from the discriminator model for scale_num in xrange(self.num_scale_nets): with tf.name_scope('scale_' + str(scale_num)): with tf.name_scope('setup'): with tf.name_scope('fully-connected'): fc_ws = [] fc_bs = [] # create weights for fc layers for i in xrange( len(self.scale_fc_layer_sizes[scale_num]) - 1): fc_ws.append( w([ self.scale_fc_layer_sizes[scale_num] [i], self.scale_fc_layer_sizes[scale_num][i + 1] ])) fc_bs.append( b([ self.scale_fc_layer_sizes[scale_num][i + 1] ])) with tf.name_scope('convolutions'): conv_ws = [] conv_bs = [] # create weights for kernels for i in xrange( len(self.scale_kernel_sizes[scale_num])): conv_ws.append( w([ self.scale_kernel_sizes[scale_num][i], self.scale_kernel_sizes[scale_num][i], self.scale_conv_layer_fms[scale_num] [i], self.scale_conv_layer_fms[scale_num][i + 1] ])) conv_bs.append( b([ self.scale_conv_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) 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(3, [inputs, last_gen_frames]) # generated frame predictions preds = inputs # perform fc multiplications with tf.name_scope('fully-connected'): for i in xrange( len(self. scale_fc_layer_sizes[scale_num]) - 1): preds = tf.nn.relu( tf.matmul(preds, fc_ws[i]) + fc_bs[i]) # reshape for convolutions preds = tf.reshape(preds, [ -1, c.FRAME_HEIGHT, c.FRAME_WIDTH, self.scale_conv_layer_fms[scale_num][0] ]) # 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, conv_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 + conv_bs[i]) else: preds = tf.nn.relu(preds + conv_bs[i]) return preds, scale_gts ## # 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 = self.scale_preds[scale_num - 1] else: last_scale_pred = None # calculate train_preds, train_gts = calculate( self.height, self.width, self.inputs, self.gt_frames, last_scale_pred) self.scale_preds.append(train_preds) self.scale_gts.append(train_gts) # 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])) ## # Training ## with tf.name_scope('train'): # global loss is the combined loss from every scale network self.global_loss = combined_loss(self.scale_preds, self.scale_gts, self.d_scale_preds) 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.scalar_summary('train_loss_G', self.global_loss) self.summaries.append(loss_summary) ## # Error ## with tf.name_scope('error'): # error computation # get error at largest scale self.psnr_error = psnr_error(self.scale_preds[-1], self.gt_frames) self.sharpdiff_error = sharp_diff_error( self.scale_preds[-1], self.gt_frames) # train error summaries summary_psnr = tf.scalar_summary('train_PSNR', self.psnr_error) summary_sharpdiff = tf.scalar_summary('train_SharpDiff', self.sharpdiff_error) self.summaries += [summary_psnr, summary_sharpdiff] # add summaries to visualize in TensorBoard self.summaries = tf.merge_summary(self.summaries)
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 for scale_num in range(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 range(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_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]) print("inputs: {}, frames: {}".format(inputs.shape, last_gen_frames.shape)) inputs = tf.concat([inputs, last_gen_frames], 3) # generated frame predictions preds = inputs # perform convolutions with tf.name_scope('convolutions'): for i in range(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 ## # 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 = 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.scale_gts_train.append(train_gts) # 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 = 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.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) ## # Error ## with tf.name_scope('error'): # error computation # get error at largest scale self.psnr_error_train = psnr_error(self.scale_preds_train[-1], self.gt_frames_train) self.sharpdiff_error_train = sharp_diff_error(self.scale_preds_train[-1], self.gt_frames_train) self.psnr_error_test = psnr_error(self.scale_preds_test[-1], self.gt_frames_test) self.sharpdiff_error_test = sharp_diff_error(self.scale_preds_test[-1], self.gt_frames_test) # train error summaries summary_psnr_train = tf.summary.scalar('train_PSNR', self.psnr_error_train) summary_sharpdiff_train = tf.summary.scalar('train_SharpDiff', self.sharpdiff_error_train) self.summaries_train += [summary_psnr_train, summary_sharpdiff_train] # test error summary_psnr_test = tf.summary.scalar('test_PSNR', self.psnr_error_test) summary_sharpdiff_test = tf.summary.scalar('test_SharpDiff', self.sharpdiff_error_test) self.summaries_test += [summary_psnr_test, summary_sharpdiff_test] # add summaries to visualize in TensorBoard self.summaries_train = tf.summary.merge(self.summaries_train) self.summaries_test = tf.summary.merge(self.summaries_test)
def define_graph(self, discriminator): """ Sets up the model graph in TensorFlow. @param discriminator: The discriminator model that discriminates frames generated by this model. """ with tf.name_scope('generator'): ## # Data ## with tf.name_scope('input'): self.input_frames_train = tf.placeholder( tf.float32, shape=[ None, self.height_train, self.width_train, 3 * c.HIST_LEN ], name='input_frames_train') self.gt_frames_train = tf.placeholder(tf.float32, shape=[ None, self.height_train, self.width_train, 3 * c.GT_LEN ], name='gt_frames_train') self.input_frames_test = tf.placeholder( tf.float32, shape=[ None, self.height_test, self.width_test, 3 * c.HIST_LEN ], name='input_frames_test') self.gt_frames_test = tf.placeholder(tf.float32, shape=[ None, self.height_test, self.width_test, 3 * c.GT_LEN ], name='gt_frames_test') # use variable batch_size for more flexibility with tf.name_scope('batch_size_train'): self.batch_size_train = tf.shape( self.input_frames_train, name='input_frames_train_shape')[0] with tf.name_scope('batch_size_test'): self.batch_size_test = tf.shape( self.input_frames_test, name='input_frames_test_shape')[0] ## # Scale network setup and calculation ## self.train_vars = [ ] # the variables to train in the optimization step 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.ws = [] self.bs = [] for scale_num in xrange(self.num_scale_nets): with tf.name_scope('scale_net_' + str(scale_num)): with tf.name_scope('setup'): scale_ws = [] scale_bs = [] # create weights for kernels with tf.name_scope('weights'): for i in xrange( len(self.scale_kernel_sizes[scale_num])): scale_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] ], 'gen_' + str(scale_num) + '_' + str(i))) with tf.name_scope('biases'): for i in xrange( len(self.scale_kernel_sizes[scale_num])): scale_bs.append( b([self.scale_layer_fms[scale_num][i + 1] ])) # add to trainable parameters self.train_vars += scale_ws self.train_vars += scale_bs self.ws.append(scale_ws) self.bs.append(scale_bs) with tf.name_scope('calculation'): with tf.name_scope('calculation_train'): ## # Perform train calculation ## if scale_num > 0: last_scale_pred_train = self.scale_preds_train[ scale_num - 1] else: last_scale_pred_train = None train_preds, train_gts = self.generate_predictions( scale_num, self.height_train, self.width_train, self.input_frames_train, self.gt_frames_train, last_scale_pred_train) with tf.name_scope('calculation_test'): ## # Perform test calculation if scale_num > 0: last_scale_pred_test = self.scale_preds_test[ scale_num - 1] else: last_scale_pred_test = None test_preds, test_gts = self.generate_predictions( scale_num, self.height_test, self.width_test, self.input_frames_test, self.gt_frames_test, last_scale_pred_test, 'test') self.scale_preds_train.append(train_preds) self.scale_gts_train.append(train_gts) self.scale_preds_test.append(test_preds) self.scale_gts_test.append(test_gts) ## # Get Discriminator Predictions ## if c.ADVERSARIAL: with tf.name_scope('d_preds'): # A list of the prediction tensors for each scale network self.d_scale_preds = [] for scale_num in xrange(self.num_scale_nets): with tf.name_scope('scale_' + str(scale_num)): with tf.name_scope('calculation'): input_scale_factor = 1. / self.scale_gt_inverse_scale_factor[ scale_num] input_scale_height = int(self.height_train * input_scale_factor) input_scale_width = int(self.width_train * input_scale_factor) scale_inputs_train = tf.image.resize_images( self.input_frames_train, [input_scale_height, input_scale_width]) # get predictions from the d scale networks self.d_scale_preds.append( discriminator.scale_nets[scale_num]. generate_all_predictions( scale_inputs_train, self.scale_preds_train[scale_num])) ## # Training ## with tf.name_scope('training'): # global loss is the combined loss from every scale network self.global_loss = temporal_combined_loss( self.scale_preds_train, self.scale_gts_train, self.d_scale_preds) with tf.name_scope('train_step'): self.global_step = tf.Variable(0, trainable=False, name='global_step') 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, var_list=self.train_vars, name='train_op') # train loss summary loss_summary = tf.summary.scalar('train_loss_G', self.global_loss) self.summaries_train.append(loss_summary) ## # Error ## with tf.name_scope('error'): # error computation # get error at largest scale with tf.name_scope('psnr_train'): self.psnr_error_train = [] for gt_num in xrange(c.GT_LEN): self.psnr_error_train.append( psnr_error( self.scale_preds_train[-1][:, :, :, gt_num * 3:(gt_num + 1) * 3], self.gt_frames_train[:, :, :, gt_num * 3:(gt_num + 1) * 3])) with tf.name_scope('sharpdiff_train'): self.sharpdiff_error_train = [] for gt_num in xrange(c.GT_LEN): self.sharpdiff_error_train.append( sharp_diff_error( self.scale_preds_train[-1][:, :, :, gt_num * 3:(gt_num + 1) * 3], self.gt_frames_train[:, :, :, gt_num * 3:(gt_num + 1) * 3])) with tf.name_scope('ssim_train'): self.ssim_error_train = [] for gt_num in xrange(c.GT_LEN): self.ssim_error_train.append( ssim_error( self.scale_preds_train[-1][:, :, :, gt_num * 3:(gt_num + 1) * 3], self.gt_frames_train[:, :, :, gt_num * 3:(gt_num + 1) * 3])) with tf.name_scope('psnr_test'): self.psnr_error_test = [] for gt_num in xrange(c.GT_LEN): self.psnr_error_test.append( psnr_error( self.scale_preds_test[-1][:, :, :, gt_num * 3:(gt_num + 1) * 3], self.gt_frames_test[:, :, :, gt_num * 3:(gt_num + 1) * 3])) with tf.name_scope('sharpdiff_test'): self.sharpdiff_error_test = [] for gt_num in xrange(c.GT_LEN): self.sharpdiff_error_test.append( sharp_diff_error( self.scale_preds_test[-1][:, :, :, gt_num * 3:(gt_num + 1) * 3], self.gt_frames_test[:, :, :, gt_num * 3:(gt_num + 1) * 3])) with tf.name_scope('ssim_test'): self.ssim_error_test = [] for gt_num in xrange(c.GT_LEN): self.ssim_error_test.append( ssim_error( self.scale_preds_test[-1][:, :, :, gt_num * 3:(gt_num + 1) * 3], self.gt_frames_test[:, :, :, gt_num * 3:(gt_num + 1) * 3])) for gt_num in xrange(c.GT_LEN): # train error summaries summary_psnr_train = tf.summary.scalar( 'train_PSNR_' + str(gt_num), self.psnr_error_train[gt_num]) summary_sharpdiff_train = tf.summary.scalar( 'train_SharpDiff_' + str(gt_num), self.sharpdiff_error_train[gt_num]) summary_ssim_train = tf.summary.scalar( 'train_SSIM_' + str(gt_num), self.ssim_error_train[gt_num]) self.summaries_train += [ summary_psnr_train, summary_sharpdiff_train, summary_ssim_train ] # test error summaries summary_psnr_test = tf.summary.scalar( 'test_PSNR_' + str(gt_num), self.psnr_error_test[gt_num]) summary_sharpdiff_test = tf.summary.scalar( 'test_SharpDiff_' + str(gt_num), self.sharpdiff_error_test[gt_num]) summary_ssim_test = tf.summary.scalar( 'test_SSIM_' + str(gt_num), self.ssim_error_test[gt_num]) self.summaries_test += [ summary_psnr_test, summary_sharpdiff_test, summary_ssim_test ] # add summaries to visualize in TensorBoard self.summaries_train = tf.summary.merge(self.summaries_train) self.summaries_test = tf.summary.merge(self.summaries_test)