def train(): if FLAGS.train_dir is None: raise ValueError('Parameter train_dir must be provided') if FLAGS.task is None: raise ValueError('Parameter task must be provided') if FLAGS.model is None: raise ValueError('Parameter model must be provided') input_config_string = config_helper.GetConfigString(FLAGS.input_config) input_config = config_helper.InputConfig(input_config_string) # Training parameters. train_config_string = config_helper.GetConfigString(FLAGS.train_config) train_config = config_helper.TrainConfig(train_config_string) batch_size = train_config.batch_size initial_learning_rate = train_config.learning_rate decay_rate = train_config.decay_rate samples_per_decay = train_config.samples_per_decay # Parameters for learning-rate decay. # The formula is decay_rate ** floor(steps / decay_steps). decay_steps = samples_per_decay / batch_size decay_steps = max(decay_steps, 1) first_code = code_loader.ReadFirstCode(input_config.data) first_code_height = ( first_code.features.feature['code_shape'].int64_list.value[0]) first_code_width = ( first_code.features.feature['code_shape'].int64_list.value[1]) max_bit_depth = ( first_code.features.feature['code_shape'].int64_list.value[2]) print('Maximum code depth: {}'.format(max_bit_depth)) with tf.Graph().as_default(): ps_ops = [ "Variable", "VariableV2", "AutoReloadVariable", "VarHandleOp" ] with tf.device( tf.train.replica_device_setter(FLAGS.ps_tasks, ps_ops=ps_ops)): codes = code_loader.LoadBinaryCode(input_config=input_config, batch_size=batch_size) if input_config.unique_code_size: print('Input code size: {} x {}'.format( first_code_height, first_code_width)) codes.set_shape([ batch_size, first_code_height, first_code_width, max_bit_depth ]) else: codes.set_shape([batch_size, None, None, max_bit_depth]) codes_effective_shape = tf.shape(codes) global_step = tf.contrib.framework.create_global_step() # Apply learning-rate decay. learning_rate = tf.train.exponential_decay( learning_rate=initial_learning_rate, global_step=global_step, decay_steps=decay_steps, decay_rate=decay_rate, staircase=True) tf.contrib.deprecated.scalar_summary('Learning Rate', learning_rate) optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate, epsilon=1.0) # Create the entropy coder model. model = model_factory.GetModelRegistry().CreateModel(FLAGS.model) model_config_string = config_helper.GetConfigString( FLAGS.model_config) model.Initialize(global_step, optimizer, model_config_string) model.BuildGraph(codes) summary_op = tf.summary.merge_all() # Verify that the model can actually be trained. if model.train_op is None: raise ValueError('Input model {} is not trainable'.format( FLAGS.model)) # We disable the summary thread run by Supervisor class by passing # summary_op=None. We still pass save_summaries_secs because it is used by # the global step counter thread. is_chief = (FLAGS.task == 0) sv = tf.train.Supervisor( logdir=FLAGS.train_dir, is_chief=is_chief, global_step=global_step, # saver=model.saver, summary_op=None, save_summaries_secs=120, save_model_secs=600, recovery_wait_secs=30) sess = sv.PrepareSession(FLAGS.master) sv.StartQueueRunners(sess) step = sess.run(global_step) print('Trainer initial step: {}.'.format(step)) # Once everything has been setup properly, save the configs. if is_chief: config_helper.SaveConfig(FLAGS.train_dir, 'input_config.json', input_config_string) config_helper.SaveConfig(FLAGS.train_dir, 'model_config.json', model_config_string) config_helper.SaveConfig(FLAGS.train_dir, 'train_config.json', train_config_string) # Train the model. next_summary_time = time.time() while not sv.ShouldStop(): feed_dict = None # Once in a while, update the summaries on the chief worker. if is_chief and next_summary_time < time.time(): summary_str = sess.run(summary_op, feed_dict=feed_dict) sv.SummaryComputed(sess, summary_str) next_summary_time = time.time() + sv.save_summaries_secs else: tf_tensors = { 'train': model.train_op, 'code_length': model.average_code_length } np_tensors = sess.run(tf_tensors, feed_dict=feed_dict) print np_tensors['code_length'] sv.Stop()
def main(_): if (FLAGS.input_codes is None or FLAGS.model is None): print('\nUsage: python entropy_coder_single.py --model=progressive ' '--model_config=model_config.json' '--iteration=15\n\n') return #if FLAGS.iteration < -1 or FLAGS.iteration > 15: # print ('\n--iteration must be between 0 and 15 inclusive, or -1 to infer ' # 'from file.\n') # return #iteration = FLAGS.iteration if not tf.gfile.Exists(FLAGS.input_codes): print('\nInput codes not found.\n') return with tf.gfile.FastGFile(FLAGS.input_codes, 'rb') as code_file: contents = code_file.read() loaded_codes = np.load(io.BytesIO(contents)) assert ['codes', 'shape'] not in loaded_codes.files loaded_shape = loaded_codes['shape'] loaded_array = loaded_codes['codes'] # Unpack and recover code shapes. unpacked_codes = np.reshape( np.unpackbits(loaded_array)[:np.prod(loaded_shape)], loaded_shape) numpy_int_codes = unpacked_codes.transpose([1, 2, 3, 0, 4]) numpy_int_codes = numpy_int_codes.reshape([ numpy_int_codes.shape[0], numpy_int_codes.shape[1], numpy_int_codes.shape[2], -1 ]) numpy_codes = numpy_int_codes.astype(np.float32) * 2.0 - 1.0 with tf.Graph().as_default() as graph: # TF tensor to hold the binary codes to losslessly compress. batch_size = 1 codes = tf.placeholder(tf.float32, shape=numpy_codes.shape) # Create the entropy coder model. global_step = None optimizer = None model = model_factory.GetModelRegistry().CreateModel(FLAGS.model) model_config_string = config_helper.GetConfigString(FLAGS.model_config) model.Initialize(global_step, optimizer, model_config_string) model.BuildGraph(codes) saver = tf.train.Saver(sharded=True, keep_checkpoint_every_n_hours=12.0) with tf.Session(graph=graph) as sess: # Initialize local variables. sess.run(tf.local_variables_initializer()) # Restore model variables. saver.restore(sess, FLAGS.checkpoint) tf_tensors = {'code_length': model.average_code_length} feed_dict = {codes: numpy_codes} np_tensors = sess.run(tf_tensors, feed_dict=feed_dict) print('Additional compression ratio: {}'.format( np_tensors['code_length']))