def gelu(x, approximate=True):
"""Gaussian Error Linear Unit.
This is a smoother version of the RELU.
Original paper: https://arxiv.org/abs/1606.08415
Args:
x: float Tensor to perform activation.
approximate: use tanh approximation
Returns:
`x` with the GELU activation applied.
"""
if approximate:
cdf = 0.5 * (1.0 + tf.tanh(
(numpy.sqrt(2 / numpy.pi) * (x + 0.044715 * tf.pow(x, 3)))))
return x * cdf
else:
# TODO: check fp16
# https://github.com/tensorflow/tensorflow/issues/25052
if x.dtype.base_dtype.name == "float16":
fp32_x = tf.cast(x, tf.float32)
else:
fp32_x = x
cdf = 0.5 * (1.0 + tf.math.erf(fp32_x / numpy.sqrt(2.0)))
if x.dtype.base_dtype.name == "float16":
return x * tf.saturate_cast(cdf, tf.float16)
return x * cdf
def run(self, *in_arrays,
return_as_list=False,
# True = return a list of NumPy arrays, False = return a single NumPy array, or a tuple if there are multiple outputs.
print_progress=False, # Print progress to the console? Useful for very large input arrays.
minibatch_size=None, # Maximum minibatch size to use, None = disable batching.
num_gpus=1, # Number of GPUs to use.
out_mul=1.0, # Multiplicative constant to apply to the output(s).
out_add=0.0, # Additive constant to apply to the output(s).
out_shrink=1, # Shrink the spatial dimensions of the output(s) by the given factor.
out_dtype=None, # Convert the output to the specified data type.
**dynamic_kwargs): # Additional keyword arguments to pass into the network construction function.
assert len(in_arrays) == self.num_inputs
num_items = in_arrays[0].shape[0]
if minibatch_size is None:
minibatch_size = num_items
key = str([list(sorted(dynamic_kwargs.items())), num_gpus, out_mul, out_add, out_shrink, out_dtype])
# Build graph.
if key not in self._run_cache:
with absolute_name_scope(self.scope + '/Run'), tf.control_dependencies(None):
in_split = list(zip(*[tf.split(x, num_gpus) for x in self.input_templates]))
out_split = []
for gpu in range(num_gpus):
with tf.device('/gpu:%d' % gpu):
out_expr = self.get_output_for(*in_split[gpu], return_as_list=True, **dynamic_kwargs)
if out_mul != 1.0:
out_expr = [x * out_mul for x in out_expr]
if out_add != 0.0:
out_expr = [x + out_add for x in out_expr]
if out_shrink > 1:
ksize = [1, 1, out_shrink, out_shrink]
out_expr = [
tf.nn.avg_pool(x, ksize=ksize, strides=ksize, padding='VALID', data_format='NCHW') for x
in out_expr]
if out_dtype is not None:
if tf.as_dtype(out_dtype).is_integer:
out_expr = [tf.round(x) for x in out_expr]
out_expr = [tf.saturate_cast(x, out_dtype) for x in out_expr]
out_split.append(out_expr)
self._run_cache[key] = [tf.concat(outputs, axis=0) for outputs in zip(*out_split)]
# Run minibatches.
out_expr = self._run_cache[key]
out_arrays = [np.empty([num_items] + shape_to_list(expr.shape)[1:], expr.dtype.name) for expr in out_expr]
for mb_begin in range(0, num_items, minibatch_size):
if print_progress:
print('\r%d / %d' % (mb_begin, num_items), end='')
mb_end = min(mb_begin + minibatch_size, num_items)
mb_in = [src[mb_begin: mb_end] for src in in_arrays]
mb_out = tf.get_default_session().run(out_expr, dict(zip(self.input_templates, mb_in)))
for dst, src in zip(out_arrays, mb_out):
dst[mb_begin: mb_end] = src
# Done.
if print_progress:
print('\r%d / %d' % (num_items, num_items))
if not return_as_list:
out_arrays = out_arrays[0] if len(out_arrays) == 1 else tuple(out_arrays)
return out_arrays
def _solarize(self, x, threshold=128):
# For each pixel in the image, select the pixel
# if the value is less than the threshold.
# Otherwise, subtract 255 from the pixel.
threshold = int((threshold / MAX_LEVEL) * 256)
threshold = tf.saturate_cast(threshold, x.dtype)
return tf.where(x < threshold, x, 255 - x)
def deepfill_model(self, checkpoint_dir):
# Line tor un with gpu
sess_config = tf.ConfigProto()
sess_config.gpu_options.allow_growth = True
sess = tf.Session(config=sess_config)
model = InpaintCAModel()
input_image_ph = tf.placeholder(tf.float32,
shape=(1, self.size[0],
self.size[1] * 2, 3))
#output = model.build_server_graph(input_image_ph, dynamic=True)
output = model.build_server_graph(self.FLAGS,
input_image_ph,
reuse=tf.AUTO_REUSE)
output = (output + 1.) * 127.5
output = tf.reverse(output, [-1])
output = tf.saturate_cast(output, tf.uint8)
vars_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
assign_ops = []
for var in vars_list:
vname = var.name
from_name = vname
var_value = tf.contrib.framework.load_variable(
checkpoint_dir, from_name)
assign_ops.append(tf.assign(var, var_value))
sess.run(assign_ops)
#print('Model loaded.')
return input_image_ph, output, sess
def _legacy_output_transform_func(*expr,
out_mul=1.0,
out_add=0.0,
out_shrink=1,
out_dtype=None):
if out_mul != 1.0:
expr = [x * out_mul for x in expr]
if out_add != 0.0:
expr = [x + out_add for x in expr]
if out_shrink > 1:
ksize = [1, 1, out_shrink, out_shrink]
expr = [
tf.nn.avg_pool(x,
ksize=ksize,
strides=ksize,
padding="VALID",
data_format="NCHW") for x in expr
]
if out_dtype is not None:
if tf.as_dtype(out_dtype).is_integer:
expr = [tf.round(x) for x in expr]
expr = [tf.saturate_cast(x, out_dtype) for x in expr]
return expr
def call(self, x):
if self.norm_type == "layernorm_L2":
epsilon = self.epsilon
dtype = x.dtype
x = tf.cast(x=x, dtype=tf.float32)
mean = tf.reduce_mean(x, axis=[-1], keepdims=True)
variance = tf.reduce_mean(tf.square(x - mean),
axis=[-1],
keepdims=True)
norm_x = (x - mean) * tf.rsqrt(variance + epsilon)
result = norm_x * self.scale + self.bias
return tf.cast(x=result, dtype=dtype)
else:
dtype = x.dtype
if dtype == tf.float16:
x = tf.cast(x, dtype=tf.float32)
mean = tf.reduce_mean(x, axis=[-1], keepdims=True)
x = x - mean
variance = tf.reduce_mean(tf.abs(x), axis=[-1], keepdims=True)
norm_x = tf.div(x, variance + self.epsilon)
y = norm_x * self.scale + self.bias
if dtype == tf.float16:
y = tf.saturate_cast(y, dtype)
return y
def __init__(self, checkpoint_dir):
model = InpaintCAModel()
sess_config = tf.ConfigProto()
sess_config.gpu_options.allow_growth = True
input_image_ph = tf.placeholder(tf.float32, shape=(1, 256, 512, 3))
output = model.build_server_graph(input_image_ph, reuse=tf.AUTO_REUSE)
output = (output + 1.) * 127.5
output = tf.reverse(output, [-1])
output = tf.saturate_cast(output, tf.uint8)
# === END OF BUILD GRAPH ===
sess = tf.Session(config=sess_config)
# load pretrained model
vars_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
assign_ops = []
for var in vars_list:
vname = var.name
from_name = vname
var_value = tf.contrib.framework.load_variable(
checkpoint_dir, from_name)
assign_ops.append(tf.assign(var, var_value))
# sess.run(assign_ops)
self.sess = sess
self.graph = output
self.placeholder = input_image_ph
self.assign_ops = assign_ops
def encode_feature_map(tensors):
"""
"""
# NOTE: arXiv:1511.04587v2, accurate image super-resolution using very deep
# convolutional networks, 3.1
#
# we use d layers where layers except the first and the last are of
# the same type: 64 filters of the size 3x3x64 where a filter
# operates on 3x3 spatial region across 64 channels (feature maps).
shape = tf.shape(tensors)
h, w, c = shape[1], shape[2], shape[3]
feature_tensors = tf.squeeze(tensors, [0])
# NOTE: split to results of 64 filters
feature_tensors = tf.split(feature_tensors, 64, axis=-1)
# NOTE: concat to build rows
feature_tensors = \
[tf.concat(feature_tensors[i:i+8], axis=1) for i in range(0, 64, 8)]
# NOTE: concat to build entire map
feature_tensors = tf.concat(feature_tensors, axis=0)
# NOTE: encode as png
png_tensor = tf.saturate_cast(feature_tensors * 127.5 + 127.5, tf.uint8)
png_tensor = tf.image.encode_png(png_tensor)
return png_tensor
def build_summaries(model):
"""
"""
# NOTE: the shape of both source_images and result_images are (N,H,W,C)
source_images = model['source_images']
result_images = model['result_images']
batch_shape = tf.shape(source_images)
n, h, w, c = batch_shape[0], batch_shape[1], batch_shape[2], batch_shape[3]
cmp_images = tf.concat([source_images, result_images], axis=2)
cmp_images = tf.reshape(cmp_images, [1, n * h, 2 * w, c])
cmp_images = 127.5 * cmp_images + 127.5
cmp_images = tf.saturate_cast(cmp_images, tf.uint8)
# NOTE: summary of training images
summary_train_images = tf.summary.image(
'train_images', cmp_images, max_outputs=1)
# NOTE: summary of validation images
summary_valid_images = tf.summary.image(
'valid_images', cmp_images, max_outputs=1)
summary_loss = tf.summary.scalar('loss', model['loss'])
return {
'loss': summary_loss,
'train_images': summary_train_images,
'valid_images': summary_valid_images,
}
def build_summaries(network):
"""
"""
# summary_loss = tf.summary.scalar('transfer loss', network['loss'])
images_c = network['image_content']
images_s = network['image_styled']
images_c = tf.slice(images_c, [0, FLAGS.padding, FLAGS.padding, 0],
[-1, 256, 256, -1])
images_s = tf.slice(images_s, [0, FLAGS.padding, FLAGS.padding, 0],
[-1, 256, 256, -1])
images_c = tf.reshape(images_c, [1, FLAGS.batch_size * 256, 256, 3])
images_s = tf.reshape(images_s, [1, FLAGS.batch_size * 256, 256, 3])
images_a = tf.concat([images_c, images_s], axis=2)
images_a = images_a * 127.5 + 127.5
# images_a = tf.add(images_a, VggNet.mean_color_bgr())
images_a = tf.reverse(images_a, [3])
images_a = tf.saturate_cast(images_a, tf.uint8)
summary_image = tf.summary.image('all', images_a, max_outputs=4)
# summary_plus = tf.summary.merge([summary_image, summary_loss])
return {
# 'summary_part': summary_loss,
'summary_plus': summary_image,
}
def __init__(self, model, batch_size, randomize_noise=False):
self.batch_size = batch_size
self.initial_dlatents = np.zeros((self.batch_size, 14, 512))
model.components.synthesis.run(self.initial_dlatents,
randomize_noise=randomize_noise,
minibatch_size=self.batch_size,
custom_inputs=[
partial(
create_variable_for_generator,
batch_size=batch_size),
partial(create_stub,
batch_size=batch_size)
],
structure='fixed')
self.sess = tf.get_default_session()
self.graph = tf.get_default_graph()
self.dlatent_variable = next(v for v in tf.global_variables()
if 'learnable_dlatents' in v.name)
self.set_dlatents(self.initial_dlatents)
self.generator_output = self.graph.get_tensor_by_name(
'G_synthesis_1/_Run/concat:0')
self.generated_image = tflib.convert_images_to_uint8(
self.generator_output, nchw_to_nhwc=True, uint8_cast=False)
self.generated_image_uint8 = tf.saturate_cast(self.generated_image,
tf.uint8)
def build_summaries(model):
"""
"""
FLAGS = tf.app.flags.FLAGS
summary_loss = tf.summary.scalar('loss', model['loss'])
# NOTE: build sr_patch + hr_patch summary
batch_size = FLAGS.batch_size
lr_patch_size = FLAGS.lr_patch_size
scaling_factor = FLAGS.scaling_factor
sr_patches = model['sr_result']
hr_patches = model['hr_target']
cmp_images = tf.concat([sr_patches, hr_patches], axis=2)
shape = [
1, batch_size * lr_patch_size * scaling_factor, scaling_factor, 3]
sub_images = tf.split(cmp_images, 2 * lr_patch_size, axis=2)
sub_images = [tf.reshape(img, shape) for img in sub_images]
cmp_images = tf.concat(sub_images, axis=2)
cmp_images = tf.saturate_cast(cmp_images * 127.5 + 127.5, tf.uint8)
summary_patches = tf.summary.image('patches', cmp_images, max_outputs=1)
return {
'summary_loss': summary_loss,
'summary_patches': summary_patches,
}
def _brightness(self, x, brightness=0.4):
''' Brightness in torchvision is implemented about multiplying the factor to image,
but tensorflow.image is just implemented about adding the factor to image.
In tensorflow.image.adjust_brightness,
For regular images, `delta` should be in the range `[0,1)`,
as it is added to the image in floating point representation,
where pixel values are in the `[0,1)` range.
adjusted = math_ops.add(
flt_image, math_ops.cast(delta, flt_image.dtype), name=name)
However in torchvision docs,
Args:
brightness (float or tuple of float (min, max)): How much to jitter brightness.
brightness_factor is chosen uniformly from [max(0, 1 - brightness), 1 + brightness]
or the given [min, max]. Should be non negative numbers.
In torchvision.transforms.functional_tensor,
return _blend(img, torch.zeros_like(img), brightness_factor)
where _blend
return brightness * img1
'''
# x = tf.image.random_brightness(x, max_delta=self.args.brightness)
x = tf.cast(x, tf.float32)
delta = tf.random.uniform(shape=[],
minval=1 - brightness,
maxval=1 + brightness,
dtype=tf.float32)
x *= delta
x = tf.saturate_cast(x, tf.uint8)
return x
def quantize_image(image):
"""
Taken from Balle's implementation
"""
image = tf.round(image * 255)
image = tf.saturate_cast(image, tf.uint8)
return image
def generative_inpainting(image, mask):
h, w, _ = image.shape
grid = 8
image = image[:h // grid * grid, :w // grid * grid, :]
mask = mask[:h // grid * grid, :w // grid * grid, :]
print('Shape of image: {}'.format(image.shape))
model = InpaintCAModel()
image = np.expand_dims(image, 0)
mask = np.expand_dims(mask, 0)
input_image = np.concatenate([image, mask], axis=2)
checkpoint_dir = 'generative_inpainting/model_logs/release_imagenet_256/'
sess_config = tf.ConfigProto()
sess_config.gpu_options.allow_growth = True
with tf.Session(config=sess_config) as sess:
input_image = tf.constant(input_image, dtype=tf.float32)
output = model.build_server_graph(input_image)
output = (output + 1.) * 127.5
output = tf.reverse(output, [-1])
output = tf.saturate_cast(output, tf.uint8)
vars_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
assign_ops = []
for var in vars_list:
vname = var.name
from_name = vname
var_value = tf.contrib.framework.load_variable(
checkpoint_dir, from_name)
assign_ops.append(tf.assign(var, var_value))
sess.run(assign_ops)
result = sess.run(output)
tf.reset_default_graph()
return result
def freeze(ckpt_path):
image = tf.placeholder(tf.float32, shape=(1, None, None, 3), name='image')
mask = tf.placeholder(tf.float32, shape=(1, None, None, 3), name='mask')
model = InpaintCAModel()
input = tf.concat([image, mask],axis=2)
output = model.build_server_graph(input)
output = (output + 1.) * 127.5
output = tf.reverse(output, [-1])
output = tf.saturate_cast(output, tf.uint8)
output = tf.add(output, 0, name='output')
init_op = tf.global_variables_initializer()
restore_saver = tf.train.Saver()
with tf.Session() as sess:
sess.run(init_op)
restore_saver.restore(sess, ckpt_path)
frozen_graph_def = tf.graph_util.convert_variables_to_constants(sess, sess.graph_def,
output_node_names=['output'])
path = os.path.dirname(ckpt_path)
with open(path + '/deepfill.pb', 'wb') as f:
f.write(frozen_graph_def.SerializeToString())
print("frozen model path: {}".format( path + '/deepfill.pb'))
def inference(path, name):
with h5py.File(''.join(['datasets/coco-256.h5']), 'r') as hf:
content_images = hf['images'].value
content_names = hf['filenames'].value
content_holder = tf.placeholder(shape=[1, 256, 256, 3],
dtype=tf.float32) # Random vector
total_batch = len(content_images)
content_iter = data_iterator(content_images, content_names, 1)
generated = model.net(content_holder)
output_format = tf.saturate_cast(generated + mean_pixel, tf.uint8)
sess = tf.Session()
sess.run(tf.initialize_all_variables())
sess.run(tf.initialize_local_variables())
saver = tf.train.Saver()
saver.restore(sess, path)
for j in tqdm.tqdm(xrange(total_batch)):
content_image, content_name = content_iter.next()
print "stylize: " + str(content_name)
content_image = np.reshape(content_image,
[1, 256, 256, 3]) - mean_pixel
output_t = sess.run(output_format,
feed_dict={content_holder: content_image})
scipy.misc.imsave(
'coco_style/%s-%s.png' % (content_name[0][5:-4], name),
output_t[0])
def translate():
"""
"""
image_path_pairs = prepare_paths()
reals = tf.placeholder(shape=[None, 256, 256, 3], dtype=tf.uint8)
flow = tf.cast(reals, dtype=tf.float32) / 127.5 - 1.0
model = build_cycle_gan(flow, flow, FLAGS.mode)
fakes = tf.saturate_cast(model['fake'] * 127.5 + 127.5, tf.uint8)
# path to checkpoint
ckpt_source_path = tf.train.latest_checkpoint(FLAGS.ckpt_dir_path)
with tf.Session() as session:
session.run(tf.global_variables_initializer())
session.run(tf.local_variables_initializer())
tf.train.Saver().restore(session, ckpt_source_path)
for i in range(0, len(image_path_pairs), FLAGS.batch_size):
path_pairs = image_path_pairs[i:i + FLAGS.batch_size]
real_images = [scipy.misc.imread(p[0]) for p in path_pairs]
fake_images = session.run(fakes, feed_dict={reals: real_images})
for idx, path in enumerate(path_pairs):
image = np.concatenate([real_images[idx], fake_images[idx]],
axis=1)
scipy.misc.imsave(path[1], image)
def complete(self):
sess_config = tf.ConfigProto()
sess_config.gpu_options.allow_growth = True
print("start")
tf.reset_default_graph()
with tf.Session(config=sess_config) as sess:
self.input_image = tf.constant(self.input_image, dtype=tf.float32)
output = self.model.build_server_graph(self.input_image)
output = (output + 1.) * 127.5
output = tf.reverse(output, [-1])
output = tf.saturate_cast(output, tf.uint8)
# load pretrained model
vars_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
assign_ops = []
for var in vars_list:
vname = var.name
from_name = vname
var_value = tf.contrib.framework.load_variable(
self.model_path, from_name)
assign_ops.append(tf.assign(var, var_value))
sess.run(assign_ops)
print('Model loaded.')
result = sess.run(output)
# cv2.imwrite(args.output, result[0][:, :, ::-1])
#cv2.imshow("result", result[0][:, :, ::-1])
#cv2.waitKey()
return result[0]
def write_png(image, filename):
# Quantize the image first
image = tf.round(image * 255)
image = tf.saturate_cast(image, tf.uint8)
image_string = tf.image.encode_png(image)
tf.io.write_file(filename, image_string)
def _random_gaussian_blur(self, x, p=.5):
if tf.less(tf.random.uniform([]), p):
sigma = tf.random.uniform([], .1, 2.)
filter_shape = 3
x = tf.expand_dims(x, axis=0)
x = tf.cast(x, tf.float32)
channels = tf.shape(x)[-1]
gaussian_kernel_x = self._get_gaussian_kernel(sigma, filter_shape)
gaussian_kernel_x = gaussian_kernel_x[None, :]
gaussian_kernel_y = self._get_gaussian_kernel(sigma, filter_shape)
gaussian_kernel_y = gaussian_kernel_y[:, None]
gaussian_kernel_2d = self._get_gaussian_kernel_2d(
gaussian_kernel_y, gaussian_kernel_x)
gaussian_kernel_2d = gaussian_kernel_2d[:, :, None, None]
gaussian_kernel_2d = tf.tile(gaussian_kernel_2d,
[1, 1, channels, 1])
x = tf.nn.depthwise_conv2d(x, gaussian_kernel_2d, (1, 1, 1, 1),
"SAME")
x = tf.squeeze(x)
x = tf.saturate_cast(x, tf.uint8)
return x
def inpaint(self, image: np.ndarray[int],
mask: np.ndarray[int]) -> np.ndarray[int]:
sess_config = tf.compat.v1.ConfigProto()
sess_config.gpu_options.allow_growth = True
input_image = self.prepare_model_input(image, mask)
# ng.get_gpus(1)
model = InpaintCAModel()
tf.compat.v1.reset_default_graph()
with tf.compat.v1.Session(config=sess_config) as sess:
input_image = tf.constant(input_image, dtype=tf.float32)
output = model.build_server_graph(self.FLAGS, input_image)
output = (output + 1.) * 127.5
output = tf.reverse(output, [-1])
output = tf.saturate_cast(output, tf.uint8)
# load pretrained model
vars_list = tf.compat.v1.get_collection(
tf.compat.v1.GraphKeys.GLOBAL_VARIABLES)
assign_ops = []
for var in vars_list:
var_value = tf.contrib.framework.load_variable(
self.checkpoint_dir, var.name)
assign_ops.append(tf.compat.v1.assign(var, var_value))
sess.run(assign_ops)
result = sess.run(output)
return result[0][:, :, ::-1]
def sharpness(image, factor):
"""Implements Sharpness function from PIL using TF ops."""
orig_im = image
image = tf.cast(image, tf.float32)
# Make image 4D for conv operation
image = tf.expand_dims(image, 0)
# SMOOTH PIL Kernel
kernel = tf.constant([[1, 1, 1], [1, 5, 1], [1, 1, 1]],
dtype=tf.float32,
shape=[3, 3, 1, 1]) / 13.
# Tile across channel dimension
kernel = tf.tile(kernel, [1, 1, 3, 1])
strides = [1, 1, 1, 1]
degenerate = tf.nn.depthwise_conv2d(image,
kernel,
strides,
padding='VALID',
dilations=[1, 1])
degenerate = tf.squeeze(tf.saturate_cast(degenerate, tf.uint8), [0])
# For the borders of the resulting image, fill in the values of the
# original image.
mask = tf.ones_like(degenerate)
padded_mask = tf.pad(mask, [[1, 1], [1, 1], [0, 0]])
padded_degenerate = tf.pad(degenerate, [[1, 1], [1, 1], [0, 0]])
result = tf.where(tf.equal(padded_mask, 1), padded_degenerate, orig_im)
# Blend the final result
return blend(result, orig_im, factor)
def version_13(cls, node, **kwargs):
tensor_dict = kwargs["tensor_dict"]
x = tensor_dict[node.inputs[0]]
y_scale = tensor_dict[node.inputs[1]]
axis = node.attrs.get("axis", 1)
x = tf.cast(x, tf.float32)
x_shape = tf_shape(x)
x_rank = len(x_shape)
y_scale_shape = tf_shape(y_scale)
y_scale_rank = len(y_scale_shape)
# Reshape process is needed for per-axis quantization
# when scale is a 1-D tensor
if y_scale_rank == 1:
shape_broadcast = list([1 for _ in range(axis)] + [x_shape[axis]] +
[1 for _ in range(axis + 1, x_rank)])
y_scale = tf.reshape(y_scale, shape_broadcast)
y = tf.divide(x, y_scale)
y = tf.round(y)
if len(node.inputs) == 3:
y_zero_point = tensor_dict[node.inputs[2]]
y_dtype = y_zero_point.dtype
y_zero_point = tf.cast(y_zero_point, tf.float32)
y_zero_point = tf.reshape(
y_zero_point,
shape_broadcast) if y_scale_rank == 1 else y_zero_point
y = tf.add(y, y_zero_point)
else: # y_zero_point default dtype = uint8
y_dtype = tf.uint8
y = tf.saturate_cast(y, y_dtype)
return [y]
def build_model(session):
"""
"""
# NOTE: restore the model
FLAGS = tf.app.flags.FLAGS
saver = tf.train.import_meta_graph(FLAGS.meta_path)
saver.restore(session, FLAGS.ckpt_path)
graph = tf.get_default_graph()
# NOTE: collect named tensors
sd_images = graph.get_tensor_by_name('sd_images:0')
sr_images = graph.get_tensor_by_name('sr_images:0')
hd_images = graph.get_tensor_by_name('hd_images:0')
# NOTE: encode super resolved image as png
img_tensor = sr_images[0]
png_tensor = tf.saturate_cast(img_tensor * 127.5 + 127.5, tf.uint8)
png_tensor = tf.image.encode_png(png_tensor)
return {
'sd_images': sd_images,
'sr_images': sr_images,
'hd_images': hd_images,
'sr_image_png': png_tensor,
'hd_sd_psnrs': tf.image.psnr(hd_images, sd_images, 2.0),
'hd_sr_psnrs': tf.image.psnr(hd_images, sr_images, 2.0),
'hd_sd_ssims': tf.image.ssim(hd_images, sd_images, 2.0),
'hd_sr_ssims': tf.image.ssim(hd_images, sr_images, 2.0),
}
def summary(images, name):
"""As a hack, saves image summaries by adding to `eval_metric_ops`."""
images = tf.saturate_cast(images * 255 + 0.5, tf.uint8)
eval_metric_ops[name] = (tf.summary.image(name,
images,
max_outputs=2),
tf.no_op())
def on_epoch_end(self, epoch, logs):
v = self.model.layers[0].variables[0].numpy()
images = self.model.layers[1](v)
images = tf.saturate_cast((images + 1.0) * 127.5, tf.uint8)
with open(f'epoch{epoch:03d}.png', 'wb') as fp:
data = tf.image.encode_png(tf.squeeze(images, axis=0)).numpy()
fp.write(data)
def get_image_from_latant_code(latent_code, file_name):
tflib.init_tf()
with dnnlib.util.open_url(STYLEGAN_MODEL_URL,
cache_dir=config.cache_dir) as f:
_G, _D, Gs = pickle.load(f)
print(latent_code.shape)
latent_code = tf.constant(latent_code)
generated_img = Gs.components.synthesis.get_output_for(
latent_code, randomize_noise=False)
generated_img = tf.transpose(generated_img, [0, 2, 3, 1])
generated_img = ((generated_img + 1) / 2) * 255
generated_img_resized_to_original = tf.image.resize_images(
generated_img,
tuple(args.input_img_size),
method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
generated_img_for_display = tf.saturate_cast(
generated_img_resized_to_original, tf.uint8)
sess = tf.get_default_session()
#tf.global_variables_initializer()
sess.run(tf.variables_initializer([]))
reconstructed_imgs = sess.run(fetches=generated_img_for_display)
print("final image shape is:", reconstructed_imgs.shape)
imageio.imwrite(os.path.join(args.restorations_dir, file_name),
reconstructed_imgs[0])
def setup(opts):
global g
global sess
global input_image
model = InpaintCAModel()
g = tf.get_default_graph()
sess = tf.Session(graph=g)
sess_config = tf.ConfigProto()
sess_config.gpu_options.allow_growth = True
input_image = tf.placeholder(tf.float32, shape=(1, 256, 256 * 2, 3))
output_image = model.build_server_graph(input_image)
output_image = (output_image + 1.) * 127.5
output_image = tf.reverse(output_image, [-1])
output_image = tf.saturate_cast(output_image, tf.uint8)
vars_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
assign_ops = []
for var in vars_list:
vname = var.name
from_name = vname
var_value = tf.contrib.framework.load_variable(opts['checkpoint_dir'],
from_name)
assign_ops.append(tf.assign(var, var_value))
sess.run(assign_ops)
print('Model loaded.')
return output_image
def resize(image, size, method):
resized = tf.image.resize(image, size, method)
if resized.dtype == tf.uint8:
return resized
else:
# deal with https://github.com/tensorflow/tensorflow/issues/10722
return tf.saturate_cast(resized, tf.uint8)
def write_png(filename, image):
"""Creates graph to write a PNG image file."""
image = tf.squeeze(image, 0)
if image.dtype.is_floating:
image = tf.round(image)
if image.dtype != tf.uint8:
image = tf.saturate_cast(image, tf.uint8)
string = tf.image.encode_png(image)
return tf.io.write_file(filename, string)
def run(self, *in_arrays,
return_as_list = False, # True = return a list of NumPy arrays, False = return a single NumPy array, or a tuple if there are multiple outputs.
print_progress = False, # Print progress to the console? Useful for very large input arrays.
minibatch_size = None, # Maximum minibatch size to use, None = disable batching.
num_gpus = 1, # Number of GPUs to use.
out_mul = 1.0, # Multiplicative constant to apply to the output(s).
out_add = 0.0, # Additive constant to apply to the output(s).
out_shrink = 1, # Shrink the spatial dimensions of the output(s) by the given factor.
out_dtype = None, # Convert the output to the specified data type.
**dynamic_kwargs): # Additional keyword arguments to pass into the network construction function.
assert len(in_arrays) == self.num_inputs
num_items = in_arrays[0].shape[0]
if minibatch_size is None:
minibatch_size = num_items
key = str([list(sorted(dynamic_kwargs.items())), num_gpus, out_mul, out_add, out_shrink, out_dtype])
# Build graph.
if key not in self._run_cache:
with absolute_name_scope(self.scope + '/Run'), tf.control_dependencies(None):
in_split = list(zip(*[tf.split(x, num_gpus) for x in self.input_templates]))
out_split = []
for gpu in range(num_gpus):
with tf.device('/gpu:%d' % gpu):
out_expr = self.get_output_for(*in_split[gpu], return_as_list=True, **dynamic_kwargs)
if out_mul != 1.0:
out_expr = [x * out_mul for x in out_expr]
if out_add != 0.0:
out_expr = [x + out_add for x in out_expr]
if out_shrink > 1:
ksize = [1, 1, out_shrink, out_shrink]
out_expr = [tf.nn.avg_pool(x, ksize=ksize, strides=ksize, padding='VALID', data_format='NCHW') for x in out_expr]
if out_dtype is not None:
if tf.as_dtype(out_dtype).is_integer:
out_expr = [tf.round(x) for x in out_expr]
out_expr = [tf.saturate_cast(x, out_dtype) for x in out_expr]
out_split.append(out_expr)
self._run_cache[key] = [tf.concat(outputs, axis=0) for outputs in zip(*out_split)]
# Run minibatches.
out_expr = self._run_cache[key]
out_arrays = [np.empty([num_items] + shape_to_list(expr.shape)[1:], expr.dtype.name) for expr in out_expr]
for mb_begin in range(0, num_items, minibatch_size):
if print_progress:
print('\r%d / %d' % (mb_begin, num_items), end='')
mb_end = min(mb_begin + minibatch_size, num_items)
mb_in = [src[mb_begin : mb_end] for src in in_arrays]
mb_out = tf.get_default_session().run(out_expr, dict(zip(self.input_templates, mb_in)))
for dst, src in zip(out_arrays, mb_out):
dst[mb_begin : mb_end] = src
# Done.
if print_progress:
print('\r%d / %d' % (num_items, num_items))
if not return_as_list:
out_arrays = out_arrays[0] if len(out_arrays) == 1 else tuple(out_arrays)
return out_arrays
def get_image_encoders(data):
"""Get a tensor to turn a batch of pixels into a batch of png encoded
bytes."""
batch_size = data.get_shape().as_list()[0]
# normalise the images
# this assumes they are all positive, this could be drastically wrong
int_data = tf.saturate_cast(data / tf.reduce_max(data, 1, keep_dims=True) *
255.0,
tf.uint8)
encoders = [tf.image.encode_png(tf.reshape(int_data[i, :], [32, 32, 1]))
for i in range(batch_size)]
return tf.pack(encoders)
def apply_ops_wrapper():
update_op = self._optimizer.apply_gradients(grads_and_vars,
global_step, name)
apply_ops = []
with tf.control_dependencies([update_op]):
for grad, var in grads_and_vars:
if var.name in self._fp32_to_fp16:
dst_var = self._fp32_to_fp16[var.name]
apply_ops.append(
tf.assign(dst_var, tf.saturate_cast(var, tf.float16)))
if apply_ops:
return tf.group(apply_ops)
return update_op
def testSaturate(self):
in_types = (tf.float32,)
out_types = tf.int8, tf.uint8, tf.int16, tf.float32
with self.test_session() as sess:
for in_type in in_types:
for out_type in out_types:
lo, hi = in_type.min, in_type.max
x = tf.constant([lo, lo + 1, lo // 2, hi // 2, hi - 1, hi], dtype=in_type)
y = tf.saturate_cast(x, dtype=out_type)
self.assertEqual(y.dtype, out_type)
x, y = sess.run([x, y])
correct = np.maximum(out_type.min, np.minimum(out_type.max, x))
self.assertAllEqual(correct, y)
def main(argv=None):
style_paths = FLAGS.STYLE_IMAGES.split(',')
style_layers = FLAGS.STYLE_LAYERS.split(',')
content_path = FLAGS.CONTENT_IMAGE
content_layers = FLAGS.CONTENT_LAYERS.split(',')
style_features_t = get_style_features(style_paths, style_layers)
res = get_content_features(content_path, 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(random if FLAGS.RANDOM_INIT else image)
net, _ = vgg.net(FLAGS.VGG_PATH, initial)
content_loss = 0
for content_features, layer in zip(content_features_t, content_layers):
layer_size = tf.size(content_features)
content_loss += tf.nn.l2_loss(net[layer] - content_features) / tf.to_float(layer_size)
content_loss = FLAGS.CONTENT_WEIGHT * content_loss / len(content_layers)
style_loss = 0
for style_gram, layer in zip(style_features_t, style_layers):
layer_size = tf.size(style_gram)
style_loss += tf.nn.l2_loss(gram(net[layer]) - style_gram) / tf.to_float(layer_size)
#style_loss += tf.sqrt(tf.reduce_sum(tf.pow(gram(net[layer]) - style_gram, 2)))
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:
sess.run(tf.initialize_all_variables())
start_time = time.time()
for step in range(FLAGS.NUM_ITERATIONS):
_, 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 optimize_loss(loss,
optimizer,
optimizer_params,
learning_rate_decay_fn,
global_step=None,
dtype=tf.float32,
gradient_noise_scale=None,
gradient_multipliers=None,
clip_gradients=None,
update_ops=None,
variables=None,
name=None,
summaries=None,
colocate_gradients_with_ops=False,
increment_global_step=True,
larc_params=None,
loss_scale=1.0,
automatic_loss_scaling=None,
on_horovod=False):
"""Given loss and parameters for optimizer, returns a training op.
Various ways of passing optimizers include:
- by string specifying the name of the optimizer. See OPTIMIZER_CLS_NAMES
for full list. E.g. `optimize_loss(..., optimizer='Adam')`.
- by function taking learning rate `Tensor` as argument and returning an
`Optimizer` instance. E.g. `optimize_loss(...,
optimizer=lambda lr: tf.train.MomentumOptimizer(lr, momentum=0.5))`.
Alternatively, if `learning_rate` is `None`, the function takes no
arguments. E.g. `optimize_loss(..., learning_rate=None,
optimizer=lambda: tf.train.MomentumOptimizer(0.5, momentum=0.5))`.
- by a subclass of `Optimizer` having a single-argument constructor
(the argument is the learning rate), such as AdamOptimizer or
AdagradOptimizer. E.g. `optimize_loss(...,
optimizer=tf.train.AdagradOptimizer)`.
- by an instance of a subclass of `Optimizer`.
E.g., `optimize_loss(..., optimizer=tf.train.AdagradOptimizer(0.5))`.
Args:
loss: Scalar `Tensor`.
global_step: Scalar int `Tensor`, step counter to update on each step
unless `increment_global_step` is `False`. If not supplied,
it will be fetched from the default graph (see
`tf.train.get_global_step` for details). If it has
not been created, no step will be incremented with each weight
update. `learning_rate_decay_fn` requires `global_step`.
learning_rate: float or `Tensor`, magnitude of update per each training
step. Can be `None`.
optimizer: string, class or optimizer instance, used as trainer.
string should be name of optimizer, like 'SGD',
'Adam', 'Adagrad'. Full list in OPTIMIZER_CLS_NAMES constant.
class should be sub-class of `tf.Optimizer` that implements
`compute_gradients` and `apply_gradients` functions.
optimizer instance should be instantiation of `tf.Optimizer`
sub-class and have `compute_gradients` and `apply_gradients`
functions.
gradient_noise_scale: float or None, adds 0-mean normal noise scaled by this
value.
gradient_multipliers: dict of variables or variable names to floats.
If present, gradients for specified
variables will be multiplied by given constant.
clip_gradients: float, callable or `None`. If float, is provided, a global
clipping is applied to prevent the norm of the gradient to exceed this
value. Alternatively, a callable can be provided e.g.: adaptive_clipping.
This callable takes a `list` of `(gradients, variables)` `tuple`s and
returns the same thing with the gradients modified.
learning_rate_decay_fn: function, takes `learning_rate` and `global_step`
`Tensor`s, returns `Tensor`.
Can be used to implement any learning rate decay
functions.
For example: `tf.train.exponential_decay`.
Ignored if `learning_rate` is not supplied.
update_ops: list of update `Operation`s to execute at each step. If `None`,
uses elements of UPDATE_OPS collection. The order of execution
between `update_ops` and `loss` is non-deterministic.
variables: list of variables to optimize or
`None` to use all trainable variables.
name: The name for this operation is used to scope operations and summaries.
summaries: List of internal quantities to visualize on tensorboard. If not
set only the loss and the learning rate will be reported. The
complete list is in OPTIMIZER_SUMMARIES.
colocate_gradients_with_ops: If True, try colocating gradients with the
corresponding op.
increment_global_step: Whether to increment `global_step`. If your model
calls `optimize_loss` multiple times per training step (e.g. to optimize
different parts of the model), use this arg to avoid incrementing
`global_step` more times than necessary.
LARC_mode: 'scale' or 'clip'
LARC_nu: If not None, LARC re-scaling will be
applied https://arxiv.org/pdf/1708.03888.pdf with nu=LARC_nu
automatic_loss_scaling: if not None, use the corresponding automatic
loss scaling algorithm. Must be one of 'Backoff'
of 'LogMax'. `dtype` must be "mixed" to use ALS.
Returns:
Training op.
Raises:
ValueError: if:
* `loss` is an invalid type or shape.
* `global_step` is an invalid type or shape.
* `learning_rate` is an invalid type or value.
* `optimizer` has the wrong type.
* `clip_gradients` is neither float nor callable.
* `learning_rate` and `learning_rate_decay_fn` are supplied, but no
`global_step` is available.
* `gradients` is empty.
"""
loss = ops.convert_to_tensor(loss)
contrib_framework.assert_scalar(loss)
if global_step is None:
global_step = tf.train.get_or_create_global_step()
else:
tf.train.assert_global_step(global_step)
with vs.variable_scope(name, "OptimizeLoss", [loss, global_step]):
# Update ops take UPDATE_OPS collection if not provided.
if update_ops is None:
update_ops = set(ops.get_collection(ops.GraphKeys.UPDATE_OPS))
# Make sure update ops are ran before computing loss.
if update_ops:
loss = control_flow_ops.with_dependencies(list(update_ops), loss)
if summaries is None:
summaries = ["learning_rate", "global_gradient_norm"]
else:
for summ in summaries:
if summ not in OPTIMIZER_SUMMARIES:
raise ValueError("Summaries should be one of [%s], you provided %s." %
(", ".join(OPTIMIZER_SUMMARIES), summ))
if global_step is None:
raise ValueError("global_step is required for learning_rate_decay_fn.")
lr = learning_rate_decay_fn(global_step)
if "learning_rate" in summaries:
summary.scalar("learning_rate", lr)
# Create optimizer, given specified parameters.
if isinstance(optimizer, six.string_types):
if lr is None:
raise ValueError("Learning rate is None, but should be specified if "
"optimizer is string (%s)." % optimizer)
if optimizer not in OPTIMIZER_CLS_NAMES:
raise ValueError(
"Optimizer name should be one of [%s], you provided %s." %
(", ".join(OPTIMIZER_CLS_NAMES), optimizer))
opt = OPTIMIZER_CLS_NAMES[optimizer](learning_rate=lr, **optimizer_params)
elif (isinstance(optimizer, type) and
issubclass(optimizer, optimizer_.Optimizer)):
if lr is None:
raise ValueError("Learning rate is None, but should be specified if "
"optimizer is class (%s)." % optimizer)
opt = optimizer(learning_rate=lr, **optimizer_params)
elif isinstance(optimizer, optimizer_.Optimizer):
opt = optimizer
elif callable(optimizer):
if lr is not None:
opt = optimizer(lr, **optimizer_params)
else:
opt = optimizer(**optimizer_params)
if not isinstance(opt, optimizer_.Optimizer):
raise ValueError("Unrecognized optimizer: function should return "
"subclass of Optimizer. Got %s." % str(opt))
else:
raise ValueError("Unrecognized optimizer: should be string, "
"subclass of Optimizer, instance of "
"subclass of Optimizer or function with one argument. "
"Got %s." % str(optimizer))
# All trainable variables, if specific variables are not specified.
if variables is None:
variables = vars_.trainable_variables()
if automatic_loss_scaling is not None:
if automatic_loss_scaling not in AutomaticLossScaler.SUPPORTED_ALGOS:
raise ValueError("Unknown automatic loss scaling algorithm: %s."
% automatic_loss_sclaing)
if dtype != "mixed":
raise ValueError("Automatic loss scaling can be used only with "
"dtype=mixed.")
loss_scale = AutomaticLossScaler(algorithm=automatic_loss_scaling)
if dtype == 'mixed':
opt = MixedPrecisionOptimizerWrapper(opt, loss_scale=loss_scale)
if on_horovod:
opt = DistributedOptimizer(opt)
# Compute gradients.
gradients = opt.compute_gradients(
loss, variables,
colocate_gradients_with_ops=colocate_gradients_with_ops,
)
# Optionally add gradient noise.
if gradient_noise_scale is not None:
gradients = _add_scaled_noise_to_gradients(gradients,
gradient_noise_scale)
# Multiply some gradients.
if gradient_multipliers is not None:
gradients = _multiply_gradients(gradients, gradient_multipliers)
if not gradients:
raise ValueError(
"Empty list of (gradient, var) pairs encountered. This is most "
"likely to be caused by an improper value of gradient_multipliers.")
if "global_gradient_norm" in summaries or "gradient_norm" in summaries:
summary.scalar(
"global_norm/gradient_norm",
clip_ops.global_norm(list(map(
lambda x: tf.cast(x, tf.float32),
list(zip(*gradients))[0])
)),
)
# Optionally clip gradients by global norm.
if clip_gradients is not None and larc_params is not None:
raise AttributeError(
"LARC and gradient norm clipping should not be used together"
)
if isinstance(clip_gradients, float):
gradients = _clip_gradients_by_norm(gradients, clip_gradients)
elif callable(clip_gradients):
gradients = clip_gradients(gradients)
elif clip_gradients is not None:
raise ValueError(
"Unknown type %s for clip_gradients" % type(clip_gradients))
# Add histograms for variables, gradients and gradient norms.
for gradient, variable in gradients:
if isinstance(gradient, ops.IndexedSlices):
grad_values = gradient.values
else:
grad_values = gradient
if isinstance(variable, ops.IndexedSlices):
var_values = variable.values
else:
var_values = variable
if grad_values is not None:
var_name = variable.name.replace(":", "_")
if "gradients" in summaries:
summary.histogram("gradients/%s" % var_name, mask_nans(grad_values))
if "gradient_norm" in summaries:
summary.scalar("gradient_norm/%s" % var_name,
clip_ops.global_norm([grad_values]))
if "variables" in summaries:
summary.histogram("variables/%s" % var_name, var_values)
if "variable_norm" in summaries:
summary.scalar("variable_norm/%s" % var_name,
clip_ops.global_norm([var_values]))
if clip_gradients is not None and ("global_gradient_norm" in summaries or
"gradient_norm" in summaries):
summary.scalar(
"global_norm/clipped_gradient_norm",
clip_ops.global_norm(list(map(
lambda x: tf.cast(x, tf.float32),
list(zip(*gradients))[0])
)),
)
# LARC gradient re-scaling
if larc_params is not None:
check_params(
config=larc_params,
required_dict={'larc_eta': float},
optional_dict={
'larc_mode': ['clip', 'scale'],
'min_update': float,
'epsilon': float
},
)
larc_eta = larc_params['larc_eta']
larc_mode = larc_params.get('larc_mode', 'clip')
min_update = larc_params.get('min_update', 1e-7)
eps = larc_params.get('epsilon', 1e-7)
for idx, (g, v) in enumerate(gradients):
var_dtype = v.dtype
v_norm = tf.norm(tensor=tf.cast(v, tf.float32), ord=2)
g_norm = tf.norm(tensor=tf.cast(g, tf.float32), ord=2)
if larc_mode == 'clip':
larc_grad_update = tf.maximum(
larc_eta * v_norm / (lr * (g_norm + eps)),
min_update,
)
if "larc_summaries" in summaries:
summary.scalar('larc_clip_on/{}'.format(v.name),
tf.cast(tf.less(larc_grad_update, 1.0), tf.int32))
larc_grad_update = tf.minimum(larc_grad_update, 1.0)
else:
larc_grad_update = tf.maximum(
larc_eta * v_norm / (g_norm + eps),
min_update,
)
larc_grad_update = tf.saturate_cast(larc_grad_update, var_dtype)
gradients[idx] = (larc_grad_update * g, v)
# adding additional summary
if "larc_summaries" in summaries:
summary.scalar('larc_grad_update/{}'.format(v.name), larc_grad_update)
summary.scalar("larc_final_lr/{}".format(v.name),
tf.cast(lr, var_dtype) * larc_grad_update)
# Create gradient updates.
grad_updates = opt.apply_gradients(
gradients,
global_step=global_step if increment_global_step else None,
name="train")
# # Ensure the train_tensor computes grad_updates.
train_tensor = control_flow_ops.with_dependencies([grad_updates], loss)
return train_tensor
def main(argv=None):
if FLAGS.CONTENT_IMAGES_PATH:
content_images = reader.image(
FLAGS.BATCH_SIZE,
FLAGS.IMAGE_SIZE,
FLAGS.CONTENT_IMAGES_PATH,
epochs=1,
shuffle=False,
crop=False)
generated_images = model.net(content_images / 255.)
output_format = tf.saturate_cast(generated_images + reader.mean_pixel, tf.uint8)
with tf.Session() as sess:
file = tf.train.latest_checkpoint(FLAGS.MODEL_PATH)
if not file:
print('Could not find trained model in {}'.format(FLAGS.MODEL_PATH))
return
print('Using model from {}'.format(file))
saver = tf.train.Saver()
saver.restore(sess, file)
coord = tf.train.Coordinator()
threads = tf.train.start_queue_runners(coord=coord)
i = 0
start_time = time.time()
try:
while not coord.should_stop():
print(i)
images_t = sess.run(output_format)
elapsed = time.time() - start_time
start_time = time.time()
print('Time for one batch: {}'.format(elapsed))
for raw_image in images_t:
i += 1
misc.imsave('out{0:04d}.png'.format(i), raw_image)
except tf.errors.OutOfRangeError:
print('Done training -- epoch limit reached')
finally:
coord.request_stop()
coord.join(threads)
return
if not os.path.exists(FLAGS.MODEL_PATH):
os.makedirs(FLAGS.MODEL_PATH)
style_paths = FLAGS.STYLE_IMAGES.split(',')
style_layers = FLAGS.STYLE_LAYERS.split(',')
content_layers = FLAGS.CONTENT_LAYERS.split(',')
style_features_t = get_style_features(style_paths, style_layers)
images = reader.image(FLAGS.BATCH_SIZE, FLAGS.IMAGE_SIZE, FLAGS.TRAIN_IMAGES_PATH)
generated = model.net(images / 255.)
net, _ = vgg.net(FLAGS.VGG_PATH, tf.concat(0, [generated, images]))
content_loss = 0
for layer in content_layers:
generated_images, content_images = tf.split(0, 2, net[layer])
size = tf.size(generated_images)
content_loss += tf.nn.l2_loss(generated_images - content_images) / tf.to_float(size)
content_loss = content_loss / len(content_layers)
style_loss = 0
for style_gram, layer in zip(style_features_t, style_layers):
generated_images, _ = tf.split(0, 2, net[layer])
size = tf.size(generated_images)
for style_image in style_gram:
style_loss += tf.nn.l2_loss(tf.reduce_sum(gram(generated_images) - style_image, 0)) / tf.to_float(size)
style_loss = style_loss / len(style_layers)
loss = FLAGS.STYLE_WEIGHT * style_loss + FLAGS.CONTENT_WEIGHT * content_loss + FLAGS.TV_WEIGHT * total_variation_loss(generated)
global_step = tf.Variable(0, name="global_step", trainable=False)
train_op = tf.train.AdamOptimizer(1e-3).minimize(loss, global_step=global_step)
output_format = tf.saturate_cast(tf.concat(0, [generated, images]) + reader.mean_pixel, tf.uint8)
with tf.Session() as sess:
saver = tf.train.Saver(tf.all_variables())
file = tf.train.latest_checkpoint(FLAGS.MODEL_PATH)
if file:
print('Restoring model from {}'.format(file))
saver.restore(sess, file)
else:
print('New model initilized')
sess.run(tf.initialize_all_variables())
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 % 100 == 0:
print(step, loss_t, elapsed_time)
output_t = sess.run(output_format)
for i, raw_image in enumerate(output_t):
misc.imsave('out{}.png'.format(i), raw_image)
if step % 10000 == 0:
saver.save(sess, FLAGS.MODEL_PATH + '/fast-style-model', global_step=step)
except tf.errors.OutOfRangeError:
print('Done training -- epoch limit reached')
finally:
coord.request_stop()
coord.join(threads)
def quantize_image(image): image = tf.round(image * 255) image = tf.saturate_cast(image, tf.uint8) return image
