Skip to content
Permalink
Browse files

Decreased buffer_size to avoid OOM issue.

  • Loading branch information...
Armour committed Jun 29, 2018
1 parent 98b7172 commit f3c7da594df9774630291b149662dba357aae45a
Showing with 8 additions and 11 deletions.
  1. +4 −5 common.py
  2. +2 −2 config.py
  3. +1 −2 test.py
  4. +1 −2 train.py
@@ -71,7 +71,7 @@ def init_model(train=True):
gray_image_yuv = rgb_to_yuv(gray_image_three_channels, "gray_image_yuv")

# Build vgg model.
with tf.name_scope("content_vgg"):
with tf.name_scope("vgg16"):
vgg.build(gray_image_three_channels)

# Predict model.
@@ -83,10 +83,9 @@ def init_model(train=True):
loss = residual_encoder.get_loss(predict_val=predict, real_val=tf.slice(color_image_yuv, [0, 0, 0, 1], [-1, -1, -1, 2], name="color_image_uv"))

# Prepare optimizer.
with tf.name_scope("optimizer"):
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops):
optimizer = tf.train.AdamOptimizer().minimize(loss, global_step=global_step)
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops):
optimizer = tf.train.AdamOptimizer().minimize(loss, global_step=global_step, name='optimizer')

# Init tensorflow summaries.
print("⏳ Init tensorflow summaries...")
@@ -17,12 +17,12 @@
image_resize_method = tf.image.ResizeMethod.BILINEAR

# Parameters for neural network.
training_iters = 2000000 # The training iterations number.
training_iters = 3000000 # The training iterations number.
batch_size = 6 # Batch size for training data.
display_step = 50 # Step interval for displaying loss and saving summary during training phase.
testing_step = 1000 # Step interval for testing and saving image during training phase.
saving_step = 10000 # Step interval for saving model during training phase.
shuffle_buffer_size = 10000
shuffle_buffer_size = 2000

# UV channel normalization parameters
u_norm_para = 0.435912
@@ -21,8 +21,7 @@
summary_hook = tf.train.SummarySaverHook(output_dir=testing_summary, save_steps=display_step, scaffold=scaffold)
checkpoint_hook = tf.train.CheckpointSaverHook(checkpoint_dir=summary_path, save_steps=saving_step, scaffold=scaffold)
num_step_hook = tf.train.StopAtStepHook(num_steps=len(file_paths))
config=tf.ConfigProto(allow_soft_placement=True, log_device_placement=True)
config.gpu_options.allow_growth = True # pylint: disable=E1101
config = tf.ConfigProto(allow_soft_placement=True, log_device_placement=True, gpu_options=(tf.GPUOptions(allow_growth=True))
session_creator = tf.train.ChiefSessionCreator(scaffold=scaffold, config=config, checkpoint_dir=summary_path)

# Create a session for running operations in the Graph.
@@ -21,8 +21,7 @@
summary_hook = tf.train.SummarySaverHook(output_dir=training_summary, save_steps=display_step, scaffold=scaffold)
checkpoint_hook = tf.train.CheckpointSaverHook(checkpoint_dir=summary_path, save_steps=saving_step, scaffold=scaffold)
num_step_hook = tf.train.StopAtStepHook(num_steps=training_iters)
config = tf.ConfigProto(allow_soft_placement=True, log_device_placement=True)
config.gpu_options.allow_growth = True # pylint: disable=E1101
config = tf.ConfigProto(allow_soft_placement=True, log_device_placement=True, gpu_options=(tf.GPUOptions(allow_growth=True))

# Create a session for running operations in the Graph.
with tf.train.MonitoredTrainingSession(checkpoint_dir=summary_path,

0 comments on commit f3c7da5

Please sign in to comment.
You can’t perform that action at this time.