def get_style_features(FLAGS):
"""
For the "style_image", the preprocessing step is:
1. Resize the shorter side to FLAGS.image_size
2. Apply central crop
"""
with tf.Graph().as_default():
network_fn = nets_factory.get_network_fn(FLAGS.loss_model,
num_classes=1,
is_training=False)
image_preprocessing_fn, image_unprocessing_fn = preprocessing_factory.get_preprocessing(
FLAGS.loss_model, is_training=False) #這裏通過配置去獲取預處理函數，與其逆函數。
# Get the style image data
size = FLAGS.image_size
img_bytes = tf.read_file(FLAGS.style_image)
if FLAGS.style_image.lower().endswith('png'):
image = tf.image.decode_png(img_bytes)
else:
image = tf.image.decode_jpeg(img_bytes)
# image = _aspect_preserving_resize(image, size)
# Add the batch dimension
images = tf.expand_dims(image_preprocessing_fn(image, size, size), 0)
# images = tf.stack([image_preprocessing_fn(image, size, size)])
_, endpoints_dict = network_fn(
images, spatial_squeeze=False) #下划线表示接受一些不需要使用的占位符。
features = []
for layer in FLAGS.style_layers:
feature = endpoints_dict[layer]
feature = tf.squeeze(gram(feature),
[0]) # remove the batch dimension
features.append(feature)
with tf.Session() as sess:
# Restore variables for loss network.
init_func = utils._get_init_fn(FLAGS)
init_func(sess)
# Make sure the 'generated' directory is exists.
if os.path.exists('generated') is False:
os.makedirs('generated')
# Indicate cropped style image path
save_file = 'generated/target_style_' + FLAGS.naming + '.jpg'
# Write preprocessed style image to indicated path
with open(save_file, 'wb') as f:
target_image = image_unprocessing_fn(images[0, :])
value = tf.image.encode_jpeg(tf.cast(target_image, tf.uint8))
f.write(sess.run(value))
tf.logging.info('Target style pattern is saved to: %s.' %
save_file)
# Return the features those layers are use for measuring style loss.
return sess.run(features)
def get_style_features(FLAGS):
"""
For the "style_image", the preprocessing step is:
1. Resize the shorter side to FLAGS.image_size
2. Apply central crop
"""
with tf.Graph().as_default():
network_fn = nets_factory.get_network_fn(
FLAGS.loss_model,
num_classes=1,
is_training=False)
image_preprocessing_fn, image_unprocessing_fn = preprocessing_factory.get_preprocessing(
FLAGS.loss_model,
is_training=False)
# Get the style image data
size = FLAGS.image_size
img_bytes = tf.read_file(FLAGS.style_image)
if FLAGS.style_image.lower().endswith('png'):
image = tf.image.decode_png(img_bytes)
else:
image = tf.image.decode_jpeg(img_bytes)
# image = _aspect_preserving_resize(image, size)
# Add the batch dimension
images = tf.expand_dims(image_preprocessing_fn(image, size, size), 0)
# images = tf.stack([image_preprocessing_fn(image, size, size)])
_, endpoints_dict = network_fn(images, spatial_squeeze=False)
features = []
for layer in FLAGS.style_layers:
feature = endpoints_dict[layer]
feature = tf.squeeze(gram(feature), [0]) # remove the batch dimension
features.append(feature)
with tf.Session() as sess:
# Restore variables for loss network.
init_func = utils._get_init_fn(FLAGS)
init_func(sess)
# Make sure the 'generated' directory is exists.
if os.path.exists('generated') is False:
os.makedirs('generated')
# Indicate cropped style image path
save_file = 'generated/target_style_' + FLAGS.naming + '.jpg'
# Write preprocessed style image to indicated path
with open(save_file, 'wb') as f:
target_image = image_unprocessing_fn(images[0, :])
value = tf.image.encode_jpeg(tf.cast(target_image, tf.uint8))
f.write(sess.run(value))
tf.logging.info('Target style pattern is saved to: %s.' % save_file)
# Return the features those layers are use for measuring style loss.
return sess.run(features)
def get_style_features(FLAGS):
"""
对于风格图片,预处理步骤:
1. Resize the shorter side to FLAGS.image_size
2. Apply central crop
"""
with tf.Graph().as_default():
network_fn = nets_factory.get_network_fn(FLAGS.loss_model,
num_classes=1,
is_training=False)
image_preprocessing_fn, image_unprocessing_fn = preprocessing_factory.get_preprocessing(
FLAGS.loss_model, is_training=False)
#获取图片风格信息
size = FLAGS.image_size
img_bytes = tf.read_file(FLAGS.style_image)
if FLAGS.style_image.lower().endswith('png'):
image = tf.image.decode_png(img_bytes)
else:
image = tf.image.decode_jpeg(img_bytes)
# image = _aspect_preserving_resize(image, size)
#增加维度
images = tf.expand_dims(image_preprocessing_fn(image, size, size), 0)
# images = tf.stack([image_preprocessing_fn(image, size, size)])
_, endpoints_dict = network_fn(images, spatial_squeeze=False)
features = []
for layer in FLAGS.style_layers:
feature = endpoints_dict[layer]
feature = tf.squeeze(gram(feature),
[0]) # remove the batch dimension
features.append(feature)
with tf.Session() as sess:
#为损失网络重新保存变量
init_func = utils._get_init_fn(FLAGS)
init_func(sess)
#确保'static/img/generated'文件夹存在
if os.path.exists('static/img/generated') is False:
os.makedirs('static/img/generated')
save_file = 'static/img/generated/target_style_' + FLAGS.naming + '.jpg'
# 写入图片
with open(save_file, 'wb') as f:
target_image = image_unprocessing_fn(images[0, :])
value = tf.image.encode_jpeg(tf.cast(target_image, tf.uint8))
f.write(sess.run(value))
tf.logging.info('Target style pattern is saved to: %s.' %
save_file)
#返回风格数据
return sess.run(features)
def get_style_features(FLAGS):
"""
For the "style_image", the preprocessing step is:
1. Resize the shorter side to FLAGS.image_size
2. Apply central crop
"""
style_layers = FLAGS.style_layers.split(',')
with tf.Graph().as_default():
network_fn = nets_factory.get_network_fn(FLAGS.loss_model,
num_classes=1,
is_training=False)
image_preprocessing_fn, image_unprocessing_fn = preprocessing_factory.get_preprocessing(
FLAGS.loss_model, is_training=False)
size = FLAGS.style_size
img_bytes = tf.read_file(FLAGS.style_image)
if FLAGS.style_image.lower().endswith('png'):
image = tf.image.decode_png(img_bytes)
else:
image = tf.image.decode_jpeg(img_bytes)
# image = _aspect_preserving_resize(image, size)
images = tf.stack([image_preprocessing_fn(image, size, size)])
_, endpoints_dict = network_fn(images, spatial_squeeze=False)
features = []
for layer in style_layers:
feature = endpoints_dict[layer]
feature = tf.squeeze(gram(feature),
[0]) # remove the batch dimension
features.append(feature)
with tf.Session() as sess:
init_func = utils._get_init_fn(FLAGS)
init_func(sess)
if os.path.exists('generated') is False:
os.makedirs('generated')
save_file = 'generated/' + FLAGS.style + "_" + str(
FLAGS.style_size) + '.jpg'
with open(save_file, 'wb') as f:
target_image = image_unprocessing_fn(images[0, :])
value = tf.image.encode_jpeg(tf.cast(target_image, tf.uint8))
f.write(sess.run(value))
tf.logging.info('Target style pattern is saved to: %s.' %
save_file)
return sess.run(features)
def get_style_features(FLAGS):
"""
对于风格图片,预处理步骤:
1. Resize the shorter side to FLAGS.image_size
2. Apply central crop
"""
with tf.Graph().as_default():
network_fn = nets_factory.get_network_fn(
FLAGS.loss_model,
num_classes=1,
is_training=False)
image_preprocessing_fn, image_unprocessing_fn = preprocessing_factory.get_preprocessing(
FLAGS.loss_model,
is_training=False)
size = FLAGS.image_size
img_bytes = tf.read_file(FLAGS.style_image)
if FLAGS.style_image.lower().endswith('png'):
image = tf.image.decode_png(img_bytes)
else:
image = tf.image.decode_jpeg(img_bytes)
# image = _aspect_preserving_resize(image, size)
images = tf.stack([image_preprocessing_fn(image, size, size)])
_, endpoints_dict = network_fn(images, spatial_squeeze=False)
features = []
for layer in FLAGS.style_layers:
feature = endpoints_dict[layer]
feature = tf.squeeze(gram(feature), [0]) # remove the batch dimension
features.append(feature)
with tf.Session() as sess:
init_func = utils._get_init_fn(FLAGS)
init_func(sess)
if os.path.exists('static/img/generated') is False:
os.makedirs('static/img/generated')
save_file = 'static/img/generated/target_style_' + FLAGS.naming + '.jpg'
with open(save_file, 'wb') as f:
target_image = image_unprocessing_fn(images[0, :])
value = tf.image.encode_jpeg(tf.cast(target_image, tf.uint8))
f.write(sess.run(value))
tf.logging.info('Target style pattern is saved to: %s.' % save_file)
return sess.run(features)
def get_style_features(style_paths, style_layers):
with tf.Graph().as_default() as g:
network_fn = nets_factory.get_network_fn(FLAGS.loss_model,
num_classes=1,
is_training=False)
image_preprocessing_fn, image_unprocessing_fn = preprocessing_factory.get_preprocessing(
FLAGS.loss_model, is_training=False)
image = tf.expand_dims(
reader.get_image(FLAGS.style_image, FLAGS.image_size,
FLAGS.image_size, image_preprocessing_fn), 0)
# image = tf.expand_dims(
# _reader.get_image(FLAGS.content_image, FLAGS.image_size), 0)
_, endpoints = network_fn(image, spatial_squeeze=False)
features = []
for layer in style_layers:
features.append(gram(endpoints[layer]))
with tf.Session() as sess:
init_func = utils._get_init_fn(FLAGS)
init_func(sess)
return sess.run(features)
def get_visual_features(FLAGS):
"""
For the "style_image", the preprocessing step is:
1. Resize the shorter side to FLAGS.image_size
2. Apply central crop
"""
with tf.Graph().as_default():
network_fn = nets_factory.get_network_fn("vgg_16",
num_classes=1,
is_training=False)
image_preprocessing_fn, image_unprocessing_fn = preprocessing_factory.get_preprocessing(
"vgg_16", is_training=False)
size = 256
img_bytes = tf.read_file(FLAGS)
image = tf.image.decode_jpeg(img_bytes, channels=3)
# image = _aspect_preserving_resize(image, size)
images = tf.stack([image_preprocessing_fn(image, size, size)])
#images= preprocessing.normalize(images, norm='l2')
_, endpoints_dict = network_fn(images, spatial_squeeze=False)
visual_layers = ["vgg_16/conv5/conv5_1"]
for layer in visual_layers:
features = endpoints_dict[layer]
with tf.Session() as sess:
init_func = utils._get_init_fn()
init_func(sess)
features = sess.run(features)
features = np.ndarray.flatten(np.array(features))
#for n in xrange(512):#change the input shape to 512*1024
# features=np.array(features)
# img_features.append(np.ndarray.flatten(features[:,:,:,n]))
return features
def main(argv=None):
# style_features_t = losses.get_style_features(FLAGS)
# Make sure the training path exists.
training_path = os.path.join(FLAGS.model_path, FLAGS.naming)
if not (os.path.exists(training_path)):
os.makedirs(training_path)
"""get features"""
style_features_t = get_style_features(FLAGS.style_image,
FLAGS.style_layers)
res = get_content_features(FLAGS.content_image, FLAGS.content_layers)
content_features_t, image_t = res[:-1], res[-1]
image = tf.constant(image_t)
random = tf.random_normal(image_t.shape)
initial = tf.Variable(image)
"""Build Network"""
network_fn = nets_factory.get_network_fn(FLAGS.loss_model,
num_classes=1,
is_training=True)
image_preprocessing_fn, image_unprocessing_fn = preprocessing_factory.get_preprocessing(
FLAGS.loss_model, is_training=False)
preprocess_content_image = tf.expand_dims(
reader.get_image(FLAGS.content_image, FLAGS.image_size,
FLAGS.image_size, image_preprocessing_fn), 0)
# preprocess_content_image = tf.expand_dims(
# _reader.get_image(FLAGS.content_image, FLAGS.image_size), 0)
# preprocess_style_image = tf.expand_dims(
# reader.get_image(FLAGS.style_image, FLAGS.image_size, FLAGS.image_size, image_preprocessing_fn), 0)
_, endpoints_dict = network_fn(preprocess_content_image,
spatial_squeeze=False)
"""build loss"""
content_loss = 0
for content_features, layer in zip(content_features_t,
FLAGS.content_layers):
layer_size = tf.size(content_features)
content_loss += tf.nn.l2_loss(
endpoints_dict[layer] - content_features) / tf.to_float(layer_size)
content_loss = FLAGS.content_weight * content_loss / len(
FLAGS.content_layers)
style_loss = 0
for style_gram, layer in zip(style_features_t, FLAGS.style_layers):
layer_size = tf.size(style_gram)
style_loss += tf.nn.l2_loss(gram(endpoints_dict[layer]) -
style_gram) / tf.to_float(layer_size)
# style_loss += (gram(endpoints_dict[layer]) - style_gram)
style_loss = FLAGS.style_weight * style_loss
tv_loss = FLAGS.tv_weight * total_variation_loss(initial)
total_loss = content_loss + style_loss + tv_loss
train_op = tf.train.AdamOptimizer(FLAGS.learning_rate).minimize(total_loss)
output_image = tf.image.encode_png(
tf.saturate_cast(tf.squeeze(initial) + reader.mean_pixel, tf.uint8))
with tf.Session() as sess:
init_func = utils._get_init_fn(FLAGS)
init_func(sess)
sess.run(tf.global_variables_initializer())
start_time = time.time()
for step in range(FLAGS.step):
_, loss_t, cl, sl = sess.run(
[train_op, total_loss, content_loss, style_loss])
elapsed = time.time() - start_time
start_time = time.time()
print(step, elapsed, loss_t, cl, sl)
image_t = sess.run(output_image)
with open('out.png', 'wb') as f:
f.write(image_t)
def main(argv=None):
content_layers = FLAGS.content_layers.split(',')
style_layers = FLAGS.style_layers.split(',')
style_layers_weights = [
float(i) for i in FLAGS.style_layers_weights.split(",")
]
#num_steps_decay = 82786 * 2 / FLAGS.batch_size
dirnames = []
for root, dirs, files in os.walk(FLAGS.data_path):
dirnames += [[os.path.join(root, name) for name in files]]
style_features_t = flow_losses.get_style_features(FLAGS)
training_path = os.path.join(FLAGS.model_path, FLAGS.naming)
if not (os.path.exists(training_path)):
os.makedirs(training_path)
with tf.Session() as sess:
"""Build Network"""
network_fn = nets_factory.get_network_fn(FLAGS.loss_model,
num_classes=1,
is_training=False)
image_preprocessing_fn, image_unprocessing_fn = preprocessing_factory.get_preprocessing(
FLAGS.loss_model, is_training=False)
image_placeholder = tf.placeholder(
tf.float32,
shape=[FLAGS.batch_size, FLAGS.image_size, FLAGS.image_size, 3])
flow_placeholder = tf.placeholder(tf.float32,
shape=[
FLAGS.batch_size - 1,
FLAGS.image_size,
FLAGS.image_size, 2
])
generated = model.net(image_placeholder, FLAGS.alpha)
processed_generated = [
image_preprocessing_fn(image, FLAGS.image_size, FLAGS.image_size)
for image in tf.unstack(generated, axis=0, num=FLAGS.batch_size)
]
processed_generated = tf.stack(processed_generated)
_, endpoints_dict = network_fn(tf.concat(
[processed_generated, image_placeholder], 0),
spatial_squeeze=False)
"""Build Losses"""
content_loss = flow_losses.content_loss(endpoints_dict, content_layers)
style_loss, style_losses = flow_losses.style_loss(
endpoints_dict, style_features_t, style_layers,
style_layers_weights)
tv_loss = flow_losses.total_variation_loss(
generated) # use the unprocessed image
flow_loss = flow_losses.flow_loss(generated, flow_placeholder,
FLAGS.batch_size, FLAGS.image_size,
FLAGS.image_size)
content_loss = FLAGS.content_weight * content_loss
style_loss = FLAGS.style_weight * style_loss
tv_loss = FLAGS.tv_weight * tv_loss
flow_loss = FLAGS.flow_weight * flow_loss
loss = style_loss + content_loss + tv_loss + flow_loss
"""Prepare to Train"""
global_step = tf.Variable(0, name="global_step", trainable=False)
variable_to_train = []
for variable in tf.trainable_variables():
if not (variable.name.startswith(FLAGS.loss_model)):
variable_to_train.append(variable)
lr = tf.train.exponential_decay(learning_rate=1e-3,
global_step=global_step,
decay_steps=100000,
decay_rate=1e-1,
staircase=True)
optimizer = tf.train.AdamOptimizer(learning_rate=lr, epsilon=1e-8)
train_op = optimizer.minimize(loss,
global_step=global_step,
var_list=variable_to_train)
#train_op = tf.train.AdamOptimizer(1e-3).minimize(loss, global_step=global_step, var_list=variable_to_train)
variables_to_restore = []
for v in tf.global_variables():
if not (v.name.startswith(FLAGS.loss_model)):
variables_to_restore.append(v)
saver = tf.train.Saver(variables_to_restore)
sess.run([
tf.global_variables_initializer(),
tf.local_variables_initializer()
])
init_func = utils._get_init_fn(FLAGS)
init_func(sess)
last_file = tf.train.latest_checkpoint(training_path)
if last_file:
print('Restoring model from {}'.format(last_file))
saver.restore(sess, last_file)
"""Start Training"""
for filenames in dirnames:
images = []
images_gray = []
for filename in sorted(filenames):
print(filename)
image_bytes = io.imread(filename)
image = image_bytes - MEAN_VALUES
image_gray = color.rgb2grey(image_bytes)
images.append(image_bytes)
images_gray.append(image_gray)
flows = []
for i in xrange(len(images_gray) - 1):
flow = cv2.calcOpticalFlowFarneback(images_gray[i],
images_gray[i + 1], 0.5, 3,
15, 3, 5, 1.2, 0)
flows.append(flow)
for i in xrange(len(images) - FLAGS.batch_size):
feed_dict = {
image_placeholder:
np.stack(images[i:i + FLAGS.batch_size]),
flow_placeholder:
np.stack(flows[i:i + FLAGS.batch_size - 1])
}
_, c_loss, s_losses, t_loss, f_loss, total_loss, step = sess.run(
[
train_op, content_loss, style_losses, tv_loss,
flow_loss, loss, global_step
],
feed_dict=feed_dict)
"""logging"""
if step % 10 == 0:
print(step, c_loss, s_losses, t_loss, f_loss, total_loss)
"""checkpoint"""
if step % 10000 == 0:
saver.save(sess,
os.path.join(training_path, 'flow-loss-model'),
global_step=step)
if step == FLAGS.max_iter:
saver.save(
sess,
os.path.join(training_path, 'flow-loss-model-done'))
print("Save flow-loss-model done!")
return
if step < FLAGS.max_iter:
saver.save(sess, os.path.join(training_path,
'flow-loss-model-done'))
print("Save flow-loss-model done")
def get_style_features(FLAGS):
"""
For the "style_image", the preprocessing step is:
1. Resize the shorter side to FLAGS.image_size
2. Apply central crop
"""
with tf.Graph().as_default():
# 定义一个网络
network_fn = nets_factory.get_network_fn(FLAGS.loss_model,
num_classes=1,
is_training=False)
# 图片预处理方法
image_preprocessing_fn, image_unprocessing_fn = preprocessing_factory.get_preprocessing(
FLAGS.loss_model, is_training=False)
width, height = FLAGS.image_width, FLAGS.image_height
# 读取图片数据
img_bytes = tf.read_file(FLAGS.style_image)
# 对图片做解码
if FLAGS.style_image.lower().endswith('png'):
image = tf.image.decode_png(img_bytes)
else:
image = tf.image.decode_jpeg(img_bytes)
# 增加一个维度变为 [1,h,w,c]，vgg网络要求喂入数据为[batch_size,h,w,c]
images = tf.expand_dims(image_preprocessing_fn(image, height, width),
0)
# 前向传播
_, endpoints_dict = network_fn(images, spatial_squeeze=False)
features = []
# 按照配置文件中的指定层特征进行提取
for layer in FLAGS.style_layers:
feature = endpoints_dict[layer]
# 因为第0维度就是图片个数，肯定是1，所以压缩成一张图片
feature = tf.squeeze(gram(feature),
[0]) # remove the batch dimension
features.append(feature)
# 保存那个style图片 我觉得也可以忽略，没有什么实际作用
with tf.Session() as sess:
# Restore variables for loss network.
init_func = utils._get_init_fn(FLAGS)
init_func(sess)
# Make sure the 'generated' directory is exists.
if os.path.exists('generated') is False:
os.makedirs('generated')
# Indicate cropped style image path
save_file = 'generated/target_style_' + FLAGS.naming + '.jpg'
# Write preprocessed style image to indicated path
with open(save_file, 'wb') as f:
target_image = image_unprocessing_fn(images[0, :])
value = tf.image.encode_jpeg(tf.cast(target_image, tf.uint8))
f.write(sess.run(value))
tf.logging.info('Target style pattern is saved to: %s.' %
save_file)
# 只是自己测试，可以删除
# save_feature_file = 'generated/target_style_feature' + FLAGS.naming + '.jpg'
# with open(save_feature_file,'wb') as f:
# # 尝试把一个特征图写入
# feature_0 = tf.cast(features[2], tf.uint8)
# feature_0 = tf.expand_dims(feature_0, 2)
# feature_img = tf.image.encode_jpeg(feature_0)
# feature_img_v = sess.run(feature_img)
# f.write(feature_img_v)
# print(feature_img_v)
# tf.logging.info('feature style pattern is saved to: %s.' % save_file)
# Return the features those layers are use for measuring style loss.
return sess.run(features)
def main(FLAGS):
style_features_t = losses.get_style_features(FLAGS)
training_path = os.path.join(FLAGS.model_path, FLAGS.naming)
if not (os.path.exists(training_path)):
os.makedirs(training_path)
with tf.Graph().as_default():
with tf.Session() as sess:
"""创建Network"""
network_fn = nets_factory.get_network_fn(
FLAGS.loss_model,
num_classes=1,
is_training=False)
image_preprocessing_fn, image_unprocessing_fn = preprocessing_factory.get_preprocessing(
FLAGS.loss_model,
is_training=False)
"""训练图片预处理"""
processed_images = reader.batch_image(FLAGS.batch_size, FLAGS.image_size, FLAGS.image_size,
'train2014/', image_preprocessing_fn, epochs=FLAGS.epoch)
generated = model.transform_network(processed_images, training=True)
processed_generated = [image_preprocessing_fn(image, FLAGS.image_size, FLAGS.image_size)
for image in tf.unstack(generated, axis=0, num=FLAGS.batch_size)
]
processed_generated = tf.stack(processed_generated)
_, endpoints_dict = network_fn(tf.concat([processed_generated, processed_images], 0), spatial_squeeze=False)
tf.logging.info('Loss network layers(You can define them in "content_layers" and "style_layers"):')
for key in endpoints_dict:
tf.logging.info(key)
"""创建 Losses"""
content_loss = losses.content_loss(endpoints_dict, FLAGS.content_layers)
style_loss, style_loss_summary = losses.style_loss(endpoints_dict, style_features_t, FLAGS.style_layers)
tv_loss = losses.total_variation_loss(generated) # use the unprocessed image
loss = FLAGS.style_weight * style_loss + FLAGS.content_weight * content_loss + FLAGS.tv_weight * tv_loss
"""准备训练"""
global_step = tf.Variable(0, name="global_step", trainable=False)
variable_to_train = []
for variable in tf.trainable_variables():
# 只训练和保存生成网络中的变量
if not (variable.name.startswith(FLAGS.loss_model)):
variable_to_train.append(variable)
"""优化"""
train_op = tf.train.AdamOptimizer(1e-3).minimize(loss, global_step=global_step, var_list=variable_to_train)
variables_to_restore = []
for v in tf.global_variables():
if not (v.name.startswith(FLAGS.loss_model)):
variables_to_restore.append(v)
saver = tf.train.Saver(variables_to_restore, write_version=tf.train.SaverDef.V1)
sess.run([tf.global_variables_initializer(), tf.local_variables_initializer()])
init_func = utils._get_init_fn(FLAGS)
init_func(sess)
last_file = tf.train.latest_checkpoint(training_path)
if last_file:
tf.logging.info('Restoring model from {}'.format(last_file))
saver.restore(sess, last_file)
"""开始训练"""
coord = tf.train.Coordinator()
threads = tf.train.start_queue_runners(coord=coord)
start_time = time.time()
try:
while not coord.should_stop():
_, loss_t, step = sess.run([train_op, loss, global_step])
elapsed_time = time.time() - start_time
start_time = time.time()
if step % 10 == 0:
tf.logging.info(
'step: %d, total Loss %f, secs/step: %f,%s' % (step, loss_t, elapsed_time, time.asctime()))
"""checkpoint"""
if step % 50 == 0:
tf.logging.info('saving check point...')
saver.save(sess, os.path.join(training_path, FLAGS.naming + '.ckpt'), global_step=step)
except tf.errors.OutOfRangeError:
saver.save(sess, os.path.join(training_path, 'fast-style-model.ckpt-done'))
tf.logging.info('Done training -- epoch limit reached')
finally:
coord.request_stop()
tf.logging.info('coordinator stop')
coord.join(threads)
def main(_):
if not os.path.isdir(FLAGS.train_dir):
os.makedirs(FLAGS.train_dir)
if not FLAGS.dataset_dir:
raise ValueError(
'You must supply the dataset directory with --dataset_dir')
if not FLAGS.aug_mode:
raise ValueError('aug_mode need to be speficied.')
if (not FLAGS.train_image_height) or (not FLAGS.train_image_width):
raise ValueError(
'The image height and width must be define explicitly.')
if FLAGS.hd_data:
if FLAGS.train_image_height != 400 or FLAGS.train_image_width != 200:
FLAGS.train_image_height, FLAGS.train_image_width = 400, 200
print("set the image size to (%d, %d)" % (400, 200))
# config and print log
config_and_print_log(FLAGS)
tf.logging.set_verbosity(tf.logging.INFO)
with tf.Graph().as_default():
#######################
# Config model_deploy #
#######################
deploy_config = model_deploy.DeploymentConfig(
num_clones=FLAGS.num_clones,
clone_on_cpu=FLAGS.clone_on_cpu,
replica_id=FLAGS.task,
num_replicas=FLAGS.worker_replicas,
num_ps_tasks=FLAGS.num_ps_tasks)
# Create global_step
with tf.device(deploy_config.variables_device()):
global_step = slim.create_global_step()
#####################################
# Select the preprocessing function #
#####################################
img_func = get_img_func()
######################
# Select the dataset #
######################
dataset = dataset_factory.DataLoader(FLAGS.model_name,
FLAGS.dataset_name,
FLAGS.dataset_dir, FLAGS.set,
FLAGS.hd_data, img_func,
FLAGS.batch_size, FLAGS.batch_k,
FLAGS.max_number_of_steps,
get_pair_type())
######################
# Select the network #
######################
network_fn = nets_factory.get_network_fn(
FLAGS.model_name,
num_classes=(dataset.num_classes - FLAGS.labels_offset),
weight_decay=FLAGS.weight_decay,
is_training=True,
sample_number=FLAGS.sample_number)
####################
# Define the model #
####################
def clone_fn(tf_batch_queue):
return build_graph(tf_batch_queue, network_fn)
clones = model_deploy.create_clones(deploy_config, clone_fn,
[dataset.tf_batch_queue])
first_clone_scope = deploy_config.clone_scope(0)
# Gather update_ops from the first clone. These contain, for example,
# the updates for the batch_norm variables created by network_fn.
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS,
first_clone_scope)
# Add summaries for end_points.
end_points = clones[0].outputs
# Add summaries for losses.
loss_dict = {}
for loss in tf.get_collection(tf.GraphKeys.LOSSES, first_clone_scope):
if loss.name == 'softmax_cross_entropy_loss/value:0':
loss_dict['clf'] = loss
elif 'softmax_cross_entropy_loss' in loss.name:
loss_dict['sample_clf_' +
str(loss.name.split('/')[0].split('_')[-1])] = loss
elif 'entropy' in loss.name:
loss_dict['entropy'] = loss
else:
raise Exception('Loss type error')
#################################
# Configure the moving averages #
#################################
if FLAGS.moving_average_decay:
moving_average_variables = slim.get_model_variables()
variable_averages = tf.train.ExponentialMovingAverage(
FLAGS.moving_average_decay, global_step)
else:
moving_average_variables, variable_averages = None, None
#########################################
# Configure the optimization procedure. #
#########################################
with tf.device(deploy_config.optimizer_device()):
learning_rate = _configure_learning_rate(dataset.num_samples,
global_step, FLAGS)
optimizer = _configure_optimizer(learning_rate)
if FLAGS.sync_replicas:
# If sync_replicas is enabled, the averaging will be done in the chief
# queue runner.
optimizer = tf.train.SyncReplicasOptimizer(
opt=optimizer,
replicas_to_aggregate=FLAGS.replicas_to_aggregate,
variable_averages=variable_averages,
variables_to_average=moving_average_variables,
replica_id=tf.constant(FLAGS.task, tf.int32, shape=()),
total_num_replicas=FLAGS.worker_replicas)
elif FLAGS.moving_average_decay:
# Update ops executed locally by trainer.
update_ops.append(
variable_averages.apply(moving_average_variables))
# Variables to train.
variables_to_train = _get_variables_to_train()
# and returns a train_tensor and summary_op
# total_loss is the sum of all LOSSES and REGULARIZATION_LOSSES in tf.GraphKeys
total_loss, clones_gradients = model_deploy.optimize_clones(
clones, optimizer, var_list=variables_to_train)
# Create gradient updates.
grad_updates = optimizer.apply_gradients(clones_gradients,
global_step=global_step)
update_ops.append(grad_updates)
update_op = tf.group(*update_ops)
train_tensor = control_flow_ops.with_dependencies([update_op],
total_loss,
name='train_op')
train_tensor_list = [train_tensor]
format_str = 'step %d, loss = %.2f'
for loss_key in sorted(loss_dict.keys()):
train_tensor_list.append(loss_dict[loss_key])
format_str += (', %s_loss = ' % loss_key + '%.8f')
format_str += ' (%.1f examples/sec; %.3f sec/batch)'
# Create a saver.
saver = tf.train.Saver(tf.global_variables(), max_to_keep=1)
checkpoint_path = os.path.join(FLAGS.train_dir, 'model.ckpt')
###########################
# Kicks off the training. #
###########################
# Build an initialization operation to run below.
init = tf.global_variables_initializer()
# Start running operations on the Graph. allow_soft_placement must be set to
# True to build towers on GPU, as some of the ops do not have GPU
# implementations.
sess = tf.Session(config=tf.ConfigProto(
allow_soft_placement=True,
log_device_placement=FLAGS.log_device_placement))
sess.run(init)
# load pretrained weights
if FLAGS.checkpoint_path is not None:
print("Load the pretrained weights")
weight_ini_fn = _get_init_fn()
weight_ini_fn(sess)
else:
print("Train from the scratch")
# Start the queue runners.
tf.train.start_queue_runners(sess=sess)
# for step in xrange(FLAGS.max_number_of_steps):
for step in xrange(FLAGS.max_number_of_steps + 1):
start_time = time.time()
loss_value_list = sess.run(train_tensor_list,
feed_dict=dataset.get_feed_dict())
duration = time.time() - start_time
# assert not np.isnan(loss_value), 'Model diverged with loss = NaN'
if step % FLAGS.log_every_n_steps == 0:
# num_examples_per_step = FLAGS.batch_size * FLAGS.num_gpus
num_examples_per_step = FLAGS.batch_size
examples_per_sec = num_examples_per_step / duration
# sec_per_batch = duration / FLAGS.num_gpus
sec_per_batch = duration
print(format_str % tuple([step] + loss_value_list +
[examples_per_sec, sec_per_batch]))
# Save the model checkpoint periodically.
# if step % FLAGS.model_snapshot_steps == 0 or (step + 1) == FLAGS.max_number_of_steps:
if step % FLAGS.model_snapshot_steps == 0:
saver.save(sess, checkpoint_path, global_step=step)
print('OK...')
def main(FLAGS):
# 得到风格特征
style_features_t = losses.get_style_features(FLAGS)
# Make sure the training path exists.
training_path = os.path.join(FLAGS.model_path, FLAGS.naming)
if not (os.path.exists(training_path)):
os.makedirs(training_path)
with tf.Graph().as_default():
with tf.Session() as sess:
"""Build Network"""
# 构造vgg网络，按照FLAGS.loss_model中的网络名字，可以在/nets/nets_factory.py 中的networks_map找到对应
network_fn = nets_factory.get_network_fn(FLAGS.loss_model,
num_classes=1,
is_training=False)
# 根据不同网络做不同的预处理
image_preprocessing_fn, image_unprocessing_fn = preprocessing_factory.get_preprocessing(
FLAGS.loss_model, is_training=False)
# 读取一个批次的数据，并且预处理
# 这里的数据你可以不用coco，可以直接给一个包含很多图片的文件夹即可
# 因为coco过于大
processed_images = reader.image(FLAGS.batch_size,
FLAGS.image_height,
FLAGS.image_width,
'F:/CASIA/train_frame/real/',
image_preprocessing_fn,
epochs=FLAGS.epoch)
# 通过生成网络，生成图片，相当于y^
generated = model.net(processed_images, training=True)
# 因为一会要把生成图片喂入到后面vgg进行计算两个损失，所以要先进行预处理
processed_generated = [
image_preprocessing_fn(image, FLAGS.image_height,
FLAGS.image_width) for image in
tf.unstack(generated, axis=0, num=FLAGS.batch_size)
]
# 因为上面是list格式，所以用tf.stack堆叠成tensor
processed_generated = tf.stack(processed_generated)
# 按照batch那一个维度，拼起来，比如原来两个是[batch_size,h,w,c]，concat后变为[2*batch_size,h,w,c]
# 这样一次前向传播把y^ 和y_c的特征都计算出来了
_, endpoints_dict = network_fn(tf.concat(
[processed_generated, processed_images], 0),
spatial_squeeze=False)
# Log the structure of loss network
tf.logging.info(
'Loss network layers(You can define them in "content_layers" and "style_layers"):'
)
for key in endpoints_dict:
tf.logging.info(key)
"""Build Losses"""
# 计算三个损失
content_loss = losses.content_loss(endpoints_dict,
FLAGS.content_layers)
style_loss, style_loss_summary = losses.style_loss(
endpoints_dict, style_features_t, FLAGS.style_layers)
tv_loss = losses.total_variation_loss(
generated) # use the unprocessed image
loss = FLAGS.style_weight * style_loss + FLAGS.content_weight * content_loss + FLAGS.tv_weight * tv_loss
# Add Summary for visualization in tensorboard.
"""Add Summary"""
# 为了tensorboard，可以忽略
tf.summary.scalar('losses/content_loss', content_loss)
tf.summary.scalar('losses/style_loss', style_loss)
tf.summary.scalar('losses/regularizer_loss', tv_loss)
tf.summary.scalar('weighted_losses/weighted_content_loss',
content_loss * FLAGS.content_weight)
tf.summary.scalar('weighted_losses/weighted_style_loss',
style_loss * FLAGS.style_weight)
tf.summary.scalar('weighted_losses/weighted_regularizer_loss',
tv_loss * FLAGS.tv_weight)
tf.summary.scalar('total_loss', loss)
for layer in FLAGS.style_layers:
tf.summary.scalar('style_losses/' + layer,
style_loss_summary[layer])
tf.summary.image('generated', generated)
# tf.image_summary('processed_generated', processed_generated) # May be better?
tf.summary.image(
'origin',
tf.stack([
image_unprocessing_fn(image) for image in tf.unstack(
processed_images, axis=0, num=FLAGS.batch_size)
]))
summary = tf.summary.merge_all()
writer = tf.summary.FileWriter(training_path)
"""Prepare to Train"""
# 步数
global_step = tf.Variable(0, name="global_step", trainable=False)
variable_to_train = []
for variable in tf.trainable_variables():
# 把非vgg网络里面的可训练变量加入variable_to_train
if not (variable.name.startswith(FLAGS.loss_model)):
variable_to_train.append(variable)
# 注意var_list
train_op = tf.train.AdamOptimizer(1e-3).minimize(
loss, global_step=global_step, var_list=variable_to_train)
variables_to_restore = []
for v in tf.global_variables():
# 把非vgg中的可存储变量加入variables_to_restore
if not (v.name.startswith(FLAGS.loss_model)):
variables_to_restore.append(v)
# 注意variables_to_restore
saver = tf.train.Saver(variables_to_restore,
write_version=tf.train.SaverDef.V1)
sess.run([
tf.global_variables_initializer(),
tf.local_variables_initializer()
])
# Restore variables for loss network.
# slim的，可以根据FLAGS里面配置把网络参数加载到sess这个会话里面
init_func = utils._get_init_fn(FLAGS)
init_func(sess)
# Restore variables for training model if the checkpoint file exists.
last_file = tf.train.latest_checkpoint(training_path)
if last_file:
tf.logging.info('Restoring model from {}'.format(last_file))
saver.restore(sess, last_file)
"""Start Training"""
coord = tf.train.Coordinator()
threads = tf.train.start_queue_runners(coord=coord)
start_time = time.time()
try:
while not coord.should_stop():
_, loss_t, step = sess.run([train_op, loss, global_step])
elapsed_time = time.time() - start_time
start_time = time.time()
"""logging"""
print(step)
if step % 10 == 0:
tf.logging.info(
'step: %d, total Loss %f, secs/step: %f' %
(step, loss_t, elapsed_time))
"""summary"""
if step % 25 == 0:
tf.logging.info('adding summary...')
summary_str = sess.run(summary)
writer.add_summary(summary_str, step)
writer.flush()
"""checkpoint"""
if step % 1000 == 0:
saver.save(sess,
os.path.join(training_path,
'fast-style-model.ckpt'),
global_step=step)
except tf.errors.OutOfRangeError:
saver.save(
sess,
os.path.join(training_path, 'fast-style-model.ckpt-done'))
tf.logging.info('Done training -- epoch limit reached')
finally:
coord.request_stop()
coord.join(threads)
def main(argv=None):
content_layers = FLAGS.content_layers.split(',')
style_layers = FLAGS.style_layers.split(',')
style_layers_weights = [
float(i) for i in FLAGS.style_layers_weights.split(",")
]
#num_steps_decay = 82786 / FLAGS.batch_size
num_steps_decay = 10000
style_features_t = losses.get_style_features(FLAGS)
training_path = os.path.join(FLAGS.model_path, FLAGS.naming)
if not (os.path.exists(training_path)):
os.makedirs(training_path)
with tf.Session() as sess:
"""Build Network"""
network_fn = nets_factory.get_network_fn(FLAGS.loss_model,
num_classes=1,
is_training=False)
image_preprocessing_fn, image_unprocessing_fn = preprocessing_factory.get_preprocessing(
FLAGS.loss_model, is_training=False)
processed_images = reader.image(FLAGS.batch_size,
FLAGS.image_size,
FLAGS.image_size,
'train2014/',
image_preprocessing_fn,
epochs=FLAGS.epoch)
generated = model.net(processed_images, FLAGS.alpha)
processed_generated = [
image_preprocessing_fn(image, FLAGS.image_size, FLAGS.image_size)
for image in tf.unstack(generated, axis=0, num=FLAGS.batch_size)
]
processed_generated = tf.stack(processed_generated)
_, endpoints_dict = network_fn(tf.concat(
[processed_generated, processed_images], 0),
spatial_squeeze=False)
"""Build Losses"""
content_loss = losses.content_loss(endpoints_dict, content_layers)
style_loss, style_losses = losses.style_loss(endpoints_dict,
style_features_t,
style_layers,
style_layers_weights)
tv_loss = losses.total_variation_loss(
generated) # use the unprocessed image
content_loss = FLAGS.content_weight * content_loss
style_loss = FLAGS.style_weight * style_loss
tv_loss = FLAGS.tv_weight * tv_loss
loss = style_loss + content_loss + tv_loss
"""Prepare to Train"""
global_step = tf.Variable(0, name="global_step", trainable=False)
variable_to_train = []
for variable in tf.trainable_variables():
if not (variable.name.startswith(FLAGS.loss_model)):
variable_to_train.append(variable)
lr = tf.train.exponential_decay(learning_rate=1e-1,
global_step=global_step,
decay_steps=num_steps_decay,
decay_rate=1e-1,
staircase=True)
optimizer = tf.train.AdamOptimizer(learning_rate=lr, epsilon=1e-8)
train_op = optimizer.minimize(loss,
global_step=global_step,
var_list=variable_to_train)
#train_op = tf.train.AdamOptimizer(1e-3).minimize(loss, global_step=global_step, var_list=variable_to_train)
variables_to_restore = []
for v in tf.global_variables():
if not (v.name.startswith(FLAGS.loss_model)):
variables_to_restore.append(v)
saver = tf.train.Saver(variables_to_restore)
sess.run([
tf.global_variables_initializer(),
tf.local_variables_initializer()
])
init_func = utils._get_init_fn(FLAGS)
init_func(sess)
last_file = tf.train.latest_checkpoint(training_path)
if last_file:
saver.restore(sess, last_file)
"""Start Training"""
coord = tf.train.Coordinator()
threads = tf.train.start_queue_runners(coord=coord)
try:
while not coord.should_stop():
_, c_loss, s_losses, t_loss, total_loss, step = sess.run([
train_op, content_loss, style_losses, tv_loss, loss,
global_step
])
"""logging"""
if step % 10 == 0:
print(step, c_loss, s_losses, t_loss, total_loss)
"""checkpoint"""
if step % 10000 == 0:
saver.save(sess,
os.path.join(training_path, 'fast-style-model'),
global_step=step)
if step == FLAGS.max_iter:
saver.save(
sess,
os.path.join(training_path, 'fast-style-model-done'))
break
except tf.errors.OutOfRangeError:
saver.save(sess,
os.path.join(training_path, 'fast-style-model-done'))
tf.logging.info('Done training -- epoch limit reached')
finally:
coord.request_stop()
coord.join(threads)
def main(FLAGS):
style_features_t = losses.get_style_features(FLAGS) #style target的Gram
#make sure the training path exists
training_path = os.path.join(FLAGS.model_path,FLAGS.naming) #model/wave/ ；用于存放训练好的模型
if not (os.path.exists(training_path)):
os.makedirs(training_path)
with tf.Graph().as_default(): #默认计算图
with tf.Session() as sess:#没有as_default(),因此，走出with 语句，sess停止执行，不能在被用
"""build loss network"""
network_fn =nets_factory.get_network_fn(FLAGS.loss_model,num_classes=1,is_training=False) #取出loss model，且该model不用训练
#对要进入loss_model的content_image,和generated_image进行preprocessing
image_preprocessing_fn,image_unpreprocessing_fn = preprocessing_factory.get_preprocessing(FLAGS.loss_model,is_training=False) #取出用于loss_model的，对image进行preprocessing和unpreprocessing的function
processed_image = reader.image(FLAGS.batch_size,FLAGS.image_size,FLAGS.image_size,'train2014/',image_preprocessing_fn,epochs=FLAGS.epoch) #这里要preprocessing的image是一个batch，为training_data
generated = model.net(processed_images,training=True) #输入“图像生成网络”的image为经过preprocessing_image，“图像生成网络”为要训练的网络
processed_generated = [image_preprocessing_fn(image,FLAGS.image_size,FLAGS.image_size) for image in tf.unstack(generated,axis=0,num=FLAGS.batch_size)]
processed_generated = tf.stack(processed_generated)
#计算generated_image和content_image进入loss_model后，更layer的output
_,endpoints_dict= network_fn(tf.concat([processed_generated,processed_images],0),spatial_squeeze=False)#endpoints_dict中存储的是2类image各个layer的值
#log the structure of loss network
tf.logging.info('loss network layers(you can define them in "content layer" and "style layer"):')
for key in endpoints_dict:
tf.logging.info(key) #屏幕输出loss_model的各个layer name
"""build losses"""
content_loss = losses.content_loss(endpoints_dict,FLAGS.content_layers)
style_loss,style_loss_summary = losses.style_loss(endpoints_dict,style_features_t,FLAGS.style_layers)
tv_loss = losses.total_variation_loss(generated)
loss = FLAGS.style_weight * style_loss + FLAGS.content_weight * content_loss + FLAGS.tv_weight * tv_loss
# Add Summary for visualization in tensorboard
"""Add Summary"""
tf.summary.scalar('losses/content_loss',content_loss)
tf.summary.scalar('losses/style_loss',style_loss)
tf.summary.scalar('losses/regularizer_loss',tv_loss)
tf.summary.scalar('weighted_losses/weighted content_loss',content_loss * FLAGS.content_weight)
tf.summary.scalar('weighted_losses/weighted style_loss',style_loss * FLAGS.style_weight)
tf.summary.scalar('weighted_losses/weighted_regularizer_loss',tv_loss * FLAGS.tv_weight)
tf.summary.scalar('total_loss',loss)
for layer in FLAGS.style_layers:
tf.summary.scalar('style_losses/' + layer,style_loss_summary[layer])
tf.summary.image('genearted',generated)
tf.summary.image('origin',tf.stack([image_unprocessing_fn(image) for image in tf.unstack(processed_images,axis=0,num=FLAGS.batch_size)]))
summary = tf.summary.merge_all()
writer = tf.summary.FileWriter(training_path)
"""prepare to train"""
global_step = tf.Variable(0,name='global_step',trainable=False)#iteration step
variable_to_train = []#需要训练的变量
for variable in tf.trainable_variables():#在图像风格迁移网络（图像生成网络+损失网络）各参数中，找需要训练的参数
if not (variable.name.startswith(FLAGS.loss_model)):
variable_to_train.append(variable)
train_op = tf.train.AdamOptimizer(1e-3).minimize(loss,global_step = global_step,var_list = variable_to_train) #需要放入sess.run()
variable_to_restore = []#在所有的全局变量中，找需要恢复默认设置的变量； 注意：local_variable指的是一些临时变量和中间变量，用于线程中，线程结束则消失
for v tf.global_variables():
if not (v.name.startswith(FLAGS.loss_model)):
variables_to_restore.append(v)
saver = tf.train.Saver(variables_to_restore,write_version=tf.train.SaverDef.V1)#利用saver.restore()恢复默认设置；这里的variable_to_restore，是需要save and restore的var_list
sess.run([tf.global_variables_initializer(),tf.local_variables_initializer()])#对全局变量和局部变量进行初始化操作：即恢复默认设置
#restore variables for loss model 恢复loss model中的参数
init_func = utils._get_init_fn(FLAGS)
init_func(sess)
#restore variables for training model if the checkpoint file exists. 如果training_model已有训练好的参数，将其载入
last_file = tf.train.latest_checkpoint(training_path)#将train_path中的model参数数据取出
if last_file:
tf.logging.info('restoringmodel from {}'.format(last_file))
saver.restore(sess,last_file) #那如果last_file不存在，就不执行restore操作吗？需要restore的参数只是图像生成网络吗？
"""start training"""
coord = tf.train.Coordinator()
threads = tf.train.start_queue_runners(coord=coord)
start_time = time.time()
try:
while not coord.should_stop():#查看线程是否停止(即：是否所有数据均运行完毕）
_,loss_t,step = sess.run([train_op,loss,global_step])
elapsed_time = time.time()
"""logging"""
#print(step)
if step % 10 == 0:
tf.logging.info('step:%d, total loss %f, secs/step: %f' % (step,loss_t,elapsed_time))
"""summary"""
if step % 25 == 0:
tf.logging.info('adding summary...')
summary_str = sess.run(summary)
writer.add_summary(summary_str,step)
writer.flush()
"""checkpoint"""
if step % 1000 == 0:
saver.save(sess,os.path.join(training_path,'fast-style-model.ckpt'),global_step=step)#保存variable_to_restore中的参数值
except tf.errors.OutOfRangeError:
saver.save(sess,os.path.join(training_path,'fast-style-model.ckpt-done'))
tf.logging.info('Done training -- epoch limit reached')
finally:
coord.request_stop()#要求停止所有线程
coord.join(threads)#将线程并入主线程，删除
def main(FLAGS):
style_features_t = losses.get_style_features(FLAGS)
# Make sure the training path exists.
training_path = os.path.join(FLAGS.model_path, FLAGS.naming)
if not(os.path.exists(training_path)):
os.makedirs(training_path)
with tf.Graph().as_default():
with tf.Session() as sess:
"""Build Network"""
network_fn = nets_factory.get_network_fn(
FLAGS.loss_model,
num_classes=1,
is_training=False)
image_preprocessing_fn, image_unprocessing_fn = preprocessing_factory.get_preprocessing(
FLAGS.loss_model,
is_training=False)
processed_images = reader.image(FLAGS.batch_size, FLAGS.image_size, FLAGS.image_size,
'train2014/', image_preprocessing_fn, epochs=FLAGS.epoch)
generated = model.net(processed_images, training=True)
processed_generated = [image_preprocessing_fn(image, FLAGS.image_size, FLAGS.image_size)
for image in tf.unstack(generated, axis=0, num=FLAGS.batch_size)
]
processed_generated = tf.stack(processed_generated)
_, endpoints_dict = network_fn(tf.concat([processed_generated, processed_images], 0), spatial_squeeze=False)
# Log the structure of loss network
tf.logging.info('Loss network layers(You can define them in "content_layers" and "style_layers"):')
for key in endpoints_dict:
tf.logging.info(key)
"""Build Losses"""
content_loss = losses.content_loss(endpoints_dict, FLAGS.content_layers)
style_loss, style_loss_summary = losses.style_loss(endpoints_dict, style_features_t, FLAGS.style_layers)
tv_loss = losses.total_variation_loss(generated) # use the unprocessed image
loss = FLAGS.style_weight * style_loss + FLAGS.content_weight * content_loss + FLAGS.tv_weight * tv_loss
# Add Summary for visualization in tensorboard.
"""Add Summary"""
tf.summary.scalar('losses/content_loss', content_loss)
tf.summary.scalar('losses/style_loss', style_loss)
tf.summary.scalar('losses/regularizer_loss', tv_loss)
tf.summary.scalar('weighted_losses/weighted_content_loss', content_loss * FLAGS.content_weight)
tf.summary.scalar('weighted_losses/weighted_style_loss', style_loss * FLAGS.style_weight)
tf.summary.scalar('weighted_losses/weighted_regularizer_loss', tv_loss * FLAGS.tv_weight)
tf.summary.scalar('total_loss', loss)
for layer in FLAGS.style_layers:
tf.summary.scalar('style_losses/' + layer, style_loss_summary[layer])
tf.summary.image('generated', generated)
# tf.image_summary('processed_generated', processed_generated) # May be better?
tf.summary.image('origin', tf.stack([
image_unprocessing_fn(image) for image in tf.unstack(processed_images, axis=0, num=FLAGS.batch_size)
]))
summary = tf.summary.merge_all()
writer = tf.summary.FileWriter(training_path)
"""Prepare to Train"""
global_step = tf.Variable(0, name="global_step", trainable=False)
variable_to_train = []
for variable in tf.trainable_variables():
if not(variable.name.startswith(FLAGS.loss_model)):
variable_to_train.append(variable)
train_op = tf.train.AdamOptimizer(1e-3).minimize(loss, global_step=global_step, var_list=variable_to_train)
variables_to_restore = []
for v in tf.global_variables():
if not(v.name.startswith(FLAGS.loss_model)):
variables_to_restore.append(v)
saver = tf.train.Saver(variables_to_restore, write_version=tf.train.SaverDef.V1)
sess.run([tf.global_variables_initializer(), tf.local_variables_initializer()])
# Restore variables for loss network.
init_func = utils._get_init_fn(FLAGS)
init_func(sess)
# Restore variables for training model if the checkpoint file exists.
last_file = tf.train.latest_checkpoint(training_path)
if last_file:
tf.logging.info('Restoring model from {}'.format(last_file))
saver.restore(sess, last_file)
"""Start Training"""
coord = tf.train.Coordinator()
threads = tf.train.start_queue_runners(coord=coord)
start_time = time.time()
try:
while not coord.should_stop():
_, loss_t, step = sess.run([train_op, loss, global_step])
elapsed_time = time.time() - start_time
start_time = time.time()
"""logging"""
# print(step)
if step % 10 == 0:
tf.logging.info('step: %d, total Loss %f, secs/step: %f' % (step, loss_t, elapsed_time))
"""summary"""
if step % 25 == 0:
tf.logging.info('adding summary...')
summary_str = sess.run(summary)
writer.add_summary(summary_str, step)
writer.flush()
"""checkpoint"""
if step % 1000 == 0:
saver.save(sess, os.path.join(training_path, 'fast-style-model.ckpt'), global_step=step)
except tf.errors.OutOfRangeError:
saver.save(sess, os.path.join(training_path, 'fast-style-model.ckpt-done'))
tf.logging.info('Done training -- epoch limit reached')
finally:
coord.request_stop()
coord.join(threads)
def main(FLAGS):
style_features_t = losses.get_style_features(FLAGS)
# Make sure the training path exists.
training_path = os.path.join(FLAGS.model_path, FLAGS.naming)
if not(os.path.exists(training_path)):
os.makedirs(training_path)
with tf.Graph().as_default():
with tf.Session() as sess:
"""Build Network"""
network_fn = nets_factory.get_network_fn(
FLAGS.loss_model,
num_classes=1,
is_training=False)
image_preprocessing_fn, image_unprocessing_fn = preprocessing_factory.get_preprocessing(
FLAGS.loss_model,
is_training=False)
processed_images = reader.image(FLAGS.batch_size, FLAGS.image_size, FLAGS.image_size,
'train2014/', image_preprocessing_fn, epochs=FLAGS.epoch)
generated = model.net(processed_images, training=True)
processed_generated = [image_preprocessing_fn(image, FLAGS.image_size, FLAGS.image_size)
for image in tf.unstack(generated, axis=0, num=FLAGS.batch_size)
]
processed_generated = tf.stack(processed_generated)
_, endpoints_dict = network_fn(tf.concat([processed_generated, processed_images], 0), spatial_squeeze=False)
# Log the structure of loss network
tf.logging.info('Loss network layers(You can define them in "content_layers" and "style_layers"):')
for key in endpoints_dict:
tf.logging.info(key)
"""Build Losses"""
content_loss = losses.content_loss(endpoints_dict, FLAGS.content_layers)
style_loss, style_loss_summary = losses.style_loss(endpoints_dict, style_features_t, FLAGS.style_layers)
tv_loss = losses.total_variation_loss(generated) # use the unprocessed image
loss = FLAGS.style_weight * style_loss + FLAGS.content_weight * content_loss + FLAGS.tv_weight * tv_loss
# Add Summary for visualization in tensorboard.
"""Add Summary"""
tf.summary.scalar('losses/content_loss', content_loss)
tf.summary.scalar('losses/style_loss', style_loss)
tf.summary.scalar('losses/regularizer_loss', tv_loss)
tf.summary.scalar('weighted_losses/weighted_content_loss', content_loss * FLAGS.content_weight)
tf.summary.scalar('weighted_losses/weighted_style_loss', style_loss * FLAGS.style_weight)
tf.summary.scalar('weighted_losses/weighted_regularizer_loss', tv_loss * FLAGS.tv_weight)
tf.summary.scalar('total_loss', loss)
for layer in FLAGS.style_layers:
tf.summary.scalar('style_losses/' + layer, style_loss_summary[layer])
tf.summary.image('generated', generated)
# tf.image_summary('processed_generated', processed_generated) # May be better?
tf.summary.image('origin', tf.stack([
image_unprocessing_fn(image) for image in tf.unstack(processed_images, axis=0, num=FLAGS.batch_size)
]))
summary = tf.summary.merge_all()
writer = tf.summary.FileWriter(training_path)
"""Prepare to Train"""
global_step = tf.Variable(0, name="global_step", trainable=False)
variable_to_train = []
for variable in tf.trainable_variables():
if not(variable.name.startswith(FLAGS.loss_model)):
variable_to_train.append(variable)
train_op = tf.train.AdamOptimizer(1e-3).minimize(loss, global_step=global_step, var_list=variable_to_train)
variables_to_restore = []
for v in tf.global_variables():
if not(v.name.startswith(FLAGS.loss_model)):
variables_to_restore.append(v)
saver = tf.train.Saver(variables_to_restore, write_version=tf.train.SaverDef.V1)
sess.run([tf.global_variables_initializer(), tf.local_variables_initializer()])
# Restore variables for loss network.
init_func = utils._get_init_fn(FLAGS)
init_func(sess)
# Restore variables for training model if the checkpoint file exists.
last_file = tf.train.latest_checkpoint(training_path)
if last_file:
tf.logging.info('Restoring model from {}'.format(last_file))
saver.restore(sess, last_file)
"""Start Training"""
coord = tf.train.Coordinator()
threads = tf.train.start_queue_runners(coord=coord)
start_time = time.time()
try:
while not coord.should_stop():
_, loss_t, step = sess.run([train_op, loss, global_step])
elapsed_time = time.time() - start_time
start_time = time.time()
"""logging"""
# print(step)
if step % 10 == 0:
tf.logging.info('step: %d, total Loss %f, secs/step: %f' % (step, loss_t, elapsed_time))
"""summary"""
if step % 25 == 0:
tf.logging.info('adding summary...')
summary_str = sess.run(summary)
writer.add_summary(summary_str, step)
writer.flush()
"""checkpoint"""
if step % 1000 == 0:
saver.save(sess, os.path.join(training_path, 'fast-style-model.ckpt'), global_step=step)
except tf.errors.OutOfRangeError:
saver.save(sess, os.path.join(training_path, 'fast-style-model.ckpt-done'))
tf.logging.info('Done training -- epoch limit reached')
finally:
coord.request_stop()
coord.join(threads)
