def _compareShapeN(self, x, use_gpu=False):
np_ans = np.array(np.shape(x))
with self.test_session(use_gpu=use_gpu) as sess:
tf_ans = tf.shape_n([x, x, x])
tf_ans_64 = tf.shape_n([x, x, x], out_type=tf.int64)
result = sess.run(tf_ans)
result_64 = sess.run(tf_ans_64)
for i in range(3):
self.assertAllEqual(np_ans, result[i])
self.assertAllEqual(np_ans, result_64[i])
self.assertShapeEqual(np_ans, tf_ans[i])
def _compareShapeN(self, x, use_gpu=False):
np_ans = np.array(np.shape(x))
with self.test_session(use_gpu=use_gpu) as sess:
tf_ans = tf.shape_n([x, x, x])
tf_ans_64 = tf.shape_n([x, x, x], out_type=tf.int64)
result = sess.run(tf_ans)
result_64 = sess.run(tf_ans_64)
for i in range(3):
self.assertAllEqual(np_ans, result[i])
self.assertAllEqual(np_ans, result_64[i])
self.assertShapeEqual(np_ans, tf_ans[i])
def _verify_compatible_image_shapes(img1, img2):
"""
Checks if two image tensors are compatible for applying SSIM or PSNR.
This function checks if two sets of images have ranks at least 3, and if the
last three dimensions match.
Args:
img1: Tensor containing the first image batch.
img2: Tensor containing the second image batch.
Returns:
A tuple containing: the first tensor shape, the second tensor shape, and a
list of control_flow_ops.Assert() ops implementing the checks.
Raises:
ValueError: When static shape check fails.
"""
shape1 = img1.get_shape().with_rank_at_least(3)
shape2 = img2.get_shape().with_rank_at_least(3)
shape1[-3:].assert_is_compatible_with(shape2[-3:])
if shape1.ndims is not None and shape2.ndims is not None:
for dim1, dim2 in zip(reversed(shape1[:-3]), reversed(shape2[:-3])):
if not (dim1 == 1 or dim2 == 1 or dim1.is_compatible_with(dim2)):
raise ValueError('Two images are not compatible: %s and %s' % (shape1, shape2))
# Now assign shape tensors.
shape1, shape2 = tf.shape_n([img1, img2])
# TODO(sjhwang): Check if shape1[:-3] and shape2[:-3] are broadcastable.
checks = []
checks.append(tf.Assert(tf.greater_equal(tf.size(shape1), 3),
[shape1, shape2], summarize=10))
checks.append(tf.Assert(tf.reduce_all(tf.equal(shape1[-3:], shape2[-3:])),
[shape1, shape2], summarize=10))
return shape1, shape2, checks
def _verify_compatible_image_shapes(img1, img2):
'''
Checks if two image tensors are compatible for applying SSIM or PSNR.
This function checks if two sets of images have ranks at least 3, and if the
last three dimensions match.
Args:
img1: Tensor containing the first image batch.
img2: Tensor containing the second image batch.
Returns:
A tuple containing: the first tensor shape, the second tensor shape, and a
list of control_flow_ops.Assert() ops implementing the checks.
Raises:
ValueError: When static shape check fails.
'''
shape1 = img1.get_shape().with_rank_at_least(3)
shape2 = img2.get_shape().with_rank_at_least(3)
shape1[-3:].assert_is_compatible_with(shape2[-3:])
if shape1.ndims is not None and shape2.ndims is not None:
for dim1, dim2 in zip(reversed(shape1[:-3]), reversed(shape2[:-3])):
if not (dim1 == 1 or dim2 == 1 or dim1.is_compatible_with(dim2)):
raise ValueError('Two images are not compatible: %s and %s' % (shape1, shape2))
# Now assign shape tensors.
shape1, shape2 = tf.shape_n([img1, img2])
# TODO(sjhwang): Check if shape1[:-3] and shape2[:-3] are broadcastable.
checks = []
checks.append(tf.Assert(tf.greater_equal(tf.size(shape1), 3),[shape1, shape2], summarize=10))
checks.append(tf.Assert(tf.reduce_all(tf.equal(shape1[-3:], shape2[-3:])),[shape1, shape2], summarize=10))
return shape1, shape2, checks
def _SSIMIndexPerChannel(
img1, img2, filter_size, filter_width, max_val=255.0):
"""Computes SSIM index between img1 and img2 per color channel.
This function matches the standard SSIM implementation found at:
https://ece.uwaterloo.ca/~z70wang/research/ssim/ssim_index.m
Details:
- To reproduce a 11x11 Gaussian filter of width 1.5 is used.
- k1 = 0.01, k2 = 0.03 as in the original paper.
Args:
img1: First RGB image batch.
img2: Second RGB image batch.
filter_size: An integer, the filter size of the Gaussian kernel used.
filter_width: A float, the filter width of the Gaussian kernel used.
max_val: the dynamic range of the images (i.e., the difference between the
maximum the and minimum allowed values).
Returns:
A pair of tensors containing batch-wise and channel-wise SSIM and
contrast-structure measure. The shape is [..., channels].
"""
filter_size = tf.constant(filter_size, dtype=tf.int32)
filter_sigma = tf.constant(filter_width, dtype=img1.dtype)
shape1, shape2 = tf.shape_n([img1, img2])
# If the images are bigger than the filter size, reduce the filter size
# to the smallest image size.
# This is consistent with the implementation in
# third_party/tensorflow_models/compression/msssim.py
# but is not present in the original algorithm.
filter_size = tf.reduce_min(
tf.concat([tf.expand_dims(filter_size, axis=0),
shape1[-3:-1],
shape2[-3:-1]],
axis=0))
kernel = FSpecialGauss(filter_size, filter_sigma)
kernel = tf.tile(kernel, multiples=[1, 1, shape1[-1], 1])
# The correct compensation factor is `1.0 - tf.reduce_sum(tf.square(kernel))`,
# but to match MATLAB implementation of MS-SSIM, we use 1.0 instead.
compensation = 1.0
def reducer(x): # pylint: disable=invalid-name
shape = tf.shape(x)
x = tf.reshape(x, shape=tf.concat([[-1], shape[-3:]], 0))
y = tf.nn.depthwise_conv2d(x, kernel, strides=[1] * 4, padding='VALID')
return tf.reshape(y, tf.concat([shape[:-3], tf.shape(y)[1:]], 0))
luminance, cs = _SSIMHelper(img1, img2, reducer, max_val, compensation)
# Average over the second and the third from the last: height, width.
axes = tf.constant([-3, -2], dtype=tf.int32)
ssim = tf.reduce_mean(luminance * cs, axes)
cs = tf.reduce_mean(cs, axes)
return ssim, cs
def generate_shape_n_pb():
with tf.compat.v1.Session(graph=tf.Graph()) as sess:
x = tf.compat.v1.placeholder(dtype="int32", shape=(2, 2))
y = tf.shape_n([1, 2], name="shape_n")
tf.io.write_graph(sess.graph,
logdir="./",
name="test_shape_n.pb",
as_text=False)
def _calculateRpnLoss(predRawScores, predBoxes, predRegressions,
predAnchors, mini_batch_size, gt_boxes, feature_h, feature_w):
"""Refactoring of code from calculateRpnLoss into another function for testing"""
num_classes = tf.shape_n(gt_boxes)[0] - 1 # subtract one since background is not a class
iou_threshold_neg = s.DEF_IOU_THRESHOLD_TRAIN_NEG
iou_threshold_pos = s.DEF_IOU_THRESHOLD_TRAIN_POS
with easy_scope(name="proposal_layer_test"), tf.device("/cpu:0"):
labeled_boxes = iou_labeler(predBoxes, gt_boxes, iou_threshold_neg, iou_threshold_pos)
# Sample boxes and raw scores for loss
posIdx, negIdx = tf.py_func(lambda x: sampleBoxes(x, num_classes, mini_batch_size),
[labeled_boxes], [tf.int32, tf.int32], stateful=False, name="sampleBoxes")
# posIdx, negIdx = tf.py_func(sampleBoxes, [labeled_boxes, num_classes, mini_batch_size],
# tf.float32, stateful=False, name="sampleBoxes")
positive_raw_scores = tf.gather(predRawScores, posIdx, axis=0,
name="positive_raw_scores")
negative_raw_scores = tf.gather(predRawScores, negIdx, axis=0,
name="negative_raw_scores")
# There is no regression loss for negative examples. For the positives, we need
# to find the gt regression from anchor to gt boxes
positive_anchors = tf.gather(predAnchors, posIdx, axis=0,
name="positive_anchors")
positive_gt_boxes = tf.gather(gt_boxes,
tf.cast(tf.gather(labeled_boxes[:, 4], posIdx), dtype=tf.int32),
name="positive_gt_boxes")
positive_gt_regs = calculateRegressions(positive_anchors, positive_gt_boxes, axis=-1)
positive_raw_regressions = tf.gather(predRegressions, posIdx, axis=0,
name="positive_raw_regressions")
# Flatten regressions before passing into the huber loss function
flat_pred_regs = tf.reshape(positive_raw_regressions, [-1])
flat_gt_regs = tf.reshape(positive_gt_regs, [-1])
reg_loss = tf.losses.huber_loss(flat_pred_regs, flat_gt_regs,
reduction=tf.losses.Reduction.NONE, delta=1.0)
reg_loss = tf.reduce_sum(reg_loss)
# Class-agnostic log loss for positive examples
# Need to create a whole bunch of [0,1]s of the right length
num_pos = tf.shape(positive_raw_scores)[0]
cls_loss_pos = tf.nn.sparse_softmax_cross_entropy_with_logits(
labels=tf.ones([num_pos], dtype=tf.int32), logits=positive_raw_scores)
cls_loss_pos = tf.reduce_sum(cls_loss_pos)
# Log-loss for the negative examples
num_neg = tf.shape(negative_raw_scores)[0]
cls_loss_neg = tf.nn.sparse_softmax_cross_entropy_with_logits(
labels=tf.zeros([num_neg], dtype=tf.int32), logits=negative_raw_scores)
cls_loss_neg = tf.reduce_sum(cls_loss_neg)
# Adding up and normalizing the losses.
reg_loss /= (feature_h * feature_w) / 10.
cls_loss = (cls_loss_pos + cls_loss_neg) / mini_batch_size
total_loss = tf.add(reg_loss, cls_loss, name="total_loss")
return total_loss
def test_shapen():
import tensorflow as tf
from dace.frontend.tensorflow import TFSession
myshape = [69, 96, 666]
num_inputs = 5
inpList = [tf.ones(myshape) for _ in range(num_inputs)]
sess_tf = tf.Session()
sess_dace = TFSession()
shapes_tf = sess_tf.run(tf.shape_n(inpList))
shapes_dace = sess_dace.run(tf.shape_n(inpList))
for dc, tf in zip(shapes_dace, shapes_tf):
try:
assert (dc == tf).all()
except (AssertionError):
print(dc)
print(tf)
def __call__(self, query, previous_alignments):
"""Score the query based on the keys and values.
Args:
query: Tensor of dtype matching `self.values` and shape #Current Decoder State
`[batch_size, query_depth]`.
previous_alignments: Tensor of dtype matching `self.values` and shape
`[batch_size, alignments_size]`
(`alignments_size` is memory's `max_time`).
Returns:
alignments: Tensor of dtype matching `self.values` and shape
`[batch_size, alignments_size]` (`alignments_size` is memory's
`max_time`).
"""
# with variable_scope.variable_scope(None, "luong_attention", [query]):
# score = _luong_score(query, self._keys, self._scale)
# alignments = self._probability_fn(score, previous_alignments)
# return alignments
encoder_parts = self.memory #batch, items, dim
# eshape=encoder_parts.get_shape().as_list()
# new_encoder_parts=encoder_parts
# l=int(self.segment_length)
# print("$$$$$$$$$$$$$$$$$$$$ ",l)
# enc=tf.expand_dims(encoder_parts,3)
# for i in range(1,l):
# #temp=tf.placeholder(tf.float32, shape=(i+1,1,1,1))
# #k = tf.fill(tf.shape(temp), 1.0)
# temp = tf.nn.conv2d(enc, self.segment_kernels[i+1], strides=[1,1,1,1], padding='VALID')
# ans=tf.squeeze(temp,3)
# ans=tf.nn.relu(ans)
# new_encoder_parts=tf.concat([new_encoder_parts,ans],1)
# encoder_parts=new_encoder_parts
decoder_parts = ops.convert_to_tensor(query, dtype=self.dtype)
decoder_parts_len = len(decoder_parts.shape.as_list())
eshape, dshape = tf.shape_n([encoder_parts, decoder_parts])
encoder_parts = tf.reshape(encoder_parts, [eshape[0]*eshape[1],eshape[2]])
encoder_parts = tf.matmul(encoder_parts, self.eW_kernel)
encoder_parts = tf.reshape(encoder_parts,[eshape[0],eshape[1],-1])
encoder_parts = tf.reshape(encoder_parts,[eshape[0]*eshape[1],-1])
encoder_parts = tf.matmul(encoder_parts,self.mult_kernel)
encoder_parts = tf.reshape(encoder_parts,[eshape[0],eshape[1],-1])
combined_states = tf.nn.tanh(encoder_parts + tf.expand_dims( tf.matmul(decoder_parts, self.dW_kernel) + self.at_bias, dim=1))
#batch, items, emb
logits = tf.reshape(tf.matmul( tf.reshape(combined_states, [ -1 , self.memory.shape[-1].value]), self.kernel), [eshape[0], eshape[1], self.vocab_size])
#batch, items, vocab
logits=logits + tf.matmul(encoder_parts,tf.expand_dims(decoder_parts,1),transpose_b=True)
#adding (d^T)(e_i)
logits=logits + self.bias
logits = tf.reduce_max(logits, 1)
# logits = 0*logits + tf.nn.xw_plus_b (decoder_parts, self.kernel1, self.bias1)
tf.print(logits,message="LOGITS: ")
return logits
def preprocess_single_image(self, image):
"""Resize image to 299 * 299."""
image = tf.image.random_flip_left_right(image)
shape, = tf.shape_n([image])
h = shape[0]
w = shape[1]
d = tf.maximum(h, w)
image = tf.image.resize_image_with_crop_or_pad(image, d, d)
image = tf.image.resize_images(image, (299, 299))
image.set_shape((299, 299, 3))
image = tf.image.per_image_standardization(image)
return image
def function_factory(model, loss, train_x, train_y, pars):
# obtain the shapes of all trainable parameters in the model
shapes = tf.shape_n(model.trainable_variables)
n_tensors = len(shapes)
# stitch and partition indexes
count = 0
stitch = []
part = []
for i, shape in enumerate(shapes):
n = np.product(shape)
stitch.append(
tf.reshape(tf.range(count, count + n, dtype=tf.int32), shape))
part.extend([i] * n)
count += n
part = tf.constant(part)
@tf.function
def assign_new_model_parameters(params_1d):
params = tf.dynamic_partition(params_1d, part, n_tensors)
for i, (shape, param) in enumerate(zip(shapes, params)):
model.trainable_variables[i].assign(tf.reshape(param, shape))
# will be returned by this factory
@tf.function
def f(params_1d):
with tf.GradientTape() as g:
assign_new_model_parameters(params_1d)
# loss_value = loss(train_y, model(train_x[0], train_x[1]))
loss_value = loss(train_x[0], train_x[1], train_y)
# calculate gradients and convert to 1D tf.Tensor
grads = g.gradient(loss_value, model.trainable_variables)
grads = tf.dynamic_stitch(stitch, grads)
error = model.get_error(pars)
# track
f.iter.assign_add(1)
if f.iter % 100 == 0 or f.iter == 1:
tf.print("Epoch:", f.iter, "loss:", loss_value, "error:", error)
return loss_value, grads
# store these information as members so we can use them outside the scope
f.iter = tf.Variable(0)
f.idx = stitch
f.part = part
f.shapes = shapes
f.assign_new_model_parameters = assign_new_model_parameters
return f
def buildFunctions(model, loss, train_x, train_y, lam):
shapes = tf.shape_n(model.trainable_variables)
n_tensors = len(shapes)
count = 0
idx = []
part = []
for i, shape in enumerate(shapes):
n = np.product(shape)
idx.append(
tf.reshape(tf.range(count, count + n, dtype=tf.int32), shape))
part.extend([i] * n)
count += n
part = tf.constant(part)
@tf.function
def assign_new_model_parameters(params_1d):
params = tf.dynamic_partition(params_1d, part, n_tensors)
for i, (shape, param) in enumerate(zip(shapes, params)):
model.trainable_variables[i].assign(tf.reshape(param, shape))
@tf.function
def f(params_1d):
with tf.GradientTape() as tape:
assign_new_model_parameters(params_1d)
loss_value = loss(
model(train_x, training=True),
train_y) + lam * tf.reduce_mean(tf.square(params_1d))
grads = tape.gradient(loss_value, model.trainable_variables)
grads = tf.dynamic_stitch(idx, grads)
f.iter.assign_add(1)
if f.iter % 50000 == 0:
tf.print("Iter:", f.iter, "loss:", loss_value)
tf.py_function(f.history.append, inp=[loss_value], Tout=[])
return loss_value, grads
f.iter = tf.Variable(0)
f.idx = idx
f.part = part
f.shapes = shapes
f.assign_new_model_parameters = assign_new_model_parameters
f.history = []
return f
def loss_function_with_model_parameters(model, loss, train_x, train_y):
"""Create a new function that assign the model parameter to the model
and evaluate its value.
Args:
model : an instance of `tf.keras.Model` or its subclasses.
loss : a function with signature loss_value = loss(pred_y, true_y).
train_x : the input part of training data.
train_y : the output part of training data.
Returns:
A function that has a signature of:
loss_value = f(model_parameters).
"""
# obtain the shapes of all trainable parameters in the model
shapes = tf.shape_n(model.trainable_variables)
count = 0
sizes = []
# Record the shape of each parameter
for shape in shapes:
n = reduce(mul, shape)
sizes.append(n)
count += n
# Function accept the parameter and evaluate model
@tf.function
def func(params):
"""A function that can be used by tfq.optimizer.rotosolve_minimize.
Args:
params [in]: a 1D tf.Tensor.
Returns:
Loss function value
"""
# update the parameters of the model
start = 0
for i, size in enumerate(sizes):
model.trainable_variables[i].assign(
tf.reshape(params[start:start + size], shape))
start += size
# evaluate the loss
loss_value = loss(model(train_x, training=True), train_y)
return loss_value
return func
def LL_gradient(x_name, x_val, param_name, param_val, y_train):
shapes = tf.shape_n(param_val)
n_tensors = len(shapes)
count = 0
idx = []
part = []
for i, shape in enumerate(shapes):
n = numpy.product(shape)
idx.append(
tf.reshape(tf.range(count, count + n, dtype=tf.int32), shape))
part.extend([i] * n)
count += n
@tf.function
def assign_new_model_parameters(params_1d):
updated_params = tf.dynamic_partition(params_1d, part, n_tensors)
for i, (shape, param) in enumerate(zip(shapes, updated_params)):
param_val[i].assign(tf.reshape(param, shape))
@tf.function
def est_grad(params_1d):
# Derive the Tensorflow gradient
with tf.GradientTape() as tape:
# Call the function to update and convert the shape of parameters
assign_new_model_parameters(params_1d)
# Estimated Choice Probability
yhat = model_fun(x_name, x_val, param_name, param_val, avail_list)
# Call the cost function
loss_value = cost_fun(y_train, yhat)
# Calculate the gradient for each parameter
estimated_grad = tape.gradient(loss_value, param_val)
grads_1dim = tf.dynamic_stitch(idx, estimated_grad)
return loss_value, grads_1dim
est_grad.idx = idx
return est_grad
def VerifyCompatibleImageShapes(img1, img2):
"""Checks if two image tensors are compatible for metric computation.
This function checks if two sets of images have ranks at least 3, and if the
last three dimensions match.
Args:
img1: The first images tensor.
img2: The second images tensor.
Returns:
A tuple of the first tensor shape, the second tensor shape, and a list of
tf.Assert() implementing the checks.
Raises:
ValueError: when static shape check fails.
"""
shape1 = img1.shape.with_rank_at_least(3)
shape2 = img2.shape.with_rank_at_least(3)
if shape1.ndims is not None and shape2.ndims is not None:
for dim1, dim2 in zip(reversed(shape1[:-3]), reversed(shape2[:-3])):
# For TF V1 compatibility.
try:
dim1 = dim1.value
dim2 = dim2.value
except AttributeError:
pass
if not (dim1 in (None, 1) or dim2 in (None, 1) or dim1 == dim2):
raise ValueError('Two images are not compatible: %s and %s' %
(shape1, shape2))
else:
raise ValueError('The two images do not have a defined shape.')
# Now assign shape tensors.
shape1, shape2 = tf.shape_n([img1, img2])
checks = []
checks.append(
tf.Assert(tf.greater_equal(tf.size(shape1), 3), [shape1, shape2],
summarize=10))
checks.append(
tf.Assert(tf.reduce_all(tf.equal(shape1[-3:], shape2[-3:])),
[shape1, shape2],
summarize=10))
return shape1, shape2, checks
def _make_1d_mapping(self):
"""
Function to make the mapping to convert to and from a 1D weight tensor back to the original weight shapes
:return: trainable_shapes, indexes, partitions
"""
trainable_shapes = tf.shape_n(self.trainable_variables)
indexes = [] # indexes to map from original shapes to 1D tensor
partitions = [] # which layer each param in 1D tensor should be mapped to
n = 0 # offset for current layer indexes
for layer, shape in enumerate(trainable_shapes):
num_values = np.product(shape) # number of trainable params in layer
# create an index for current layer in current layer shape
indexes.append(tf.reshape(tf.range(n, n + num_values), shape))
# Update the list of partitions
partitions.extend([layer] * num_values)
n += num_values
return trainable_shapes, indexes, partitions
def _conv_add_relu_grad(op, grad):
act = op.get_attr('act')
y = op.outputs[0]
if act == 'RELU':
grad = gen_nn_ops.relu_grad(grad, y)
elif act == 'RELU6':
grad = gen_nn_ops.relu6_grad(grad, y)
elif act == 'TANH':
y = math_ops.conj(y)
grad = gen_math_ops.tanh_grad(y, grad)
broadcast_shape = tf.shape(y)
input_value_shape = tf.shape(op.inputs[2])
_, reduction_axes = tf.raw_ops.BroadcastGradientArgs(
s0=broadcast_shape, s1=input_value_shape)
updates_grad_reshaped = tf.reduce_sum(
grad, axis=reduction_axes, keepdims=True)
bias_grad = tf.reshape(updates_grad_reshaped, input_value_shape)
dilations = [1, op.get_attr('dilation_w'), op.get_attr('dilation_h'), 1]
strides = [1, op.get_attr('stride_w'), op.get_attr('stride_h'), 1]
padding = op.get_attr('padding')
shape_0, shape_1 = tf.shape_n([op.inputs[0], op.inputs[1]])
return [
tf.compat.v1.nn.conv2d_backprop_input(
shape_0,
op.inputs[1],
grad,
strides=strides,
padding=padding,
dilations=dilations,
data_format='NHWC'),
tf.compat.v1.nn.conv2d_backprop_filter(
op.inputs[0],
shape_1,
grad,
strides=strides,
padding=padding,
dilations=dilations,
data_format='NHWC'), bias_grad
]
def get_context(self, query, token):
encoder_parts = self.memory #batch, items, dim
print("Token ",token, self.kernel)
decoder_parts = ops.convert_to_tensor(query, dtype=self.dtype)
decoder_parts_len = len(decoder_parts.shape.as_list())
eshape, dshape = tf.shape_n([encoder_parts, decoder_parts])
encoder_parts = tf.reshape(encoder_parts, [eshape[0]*eshape[1],eshape[2]])
encoder_parts = tf.matmul(encoder_parts, self.eW_kernel)
encoder_parts = tf.reshape(encoder_parts,[eshape[0],eshape[1],-1])
encoder_parts = tf.reshape(encoder_parts,[eshape[0]*eshape[1],-1])
encoder_parts = tf.matmul(encoder_parts,self.mult_kernel)
encoder_parts = tf.reshape(encoder_parts,[eshape[0],eshape[1],-1])
print(decoder_parts_len, decoder_parts)
combined_states = tf.nn.tanh(encoder_parts + tf.expand_dims( tf.matmul(decoder_parts, self.dW_kernel) + self.at_bias, dim=1))
#batch, items, emb
word_embedding=tf.nn.embedding_lookup( tf.transpose(self.kernel),token)
print("Word Embedding ",word_embedding)
scores = tf.squeeze(tf.matmul ( combined_states, tf.expand_dims(word_embedding, dim=1), transpose_b=True), axis=2)
#batch, items
attention_values=tf.nn.softmax(scores) #batch, items
context=tf.reduce_sum(tf.expand_dims(attention_values,dim=2)*self.memory, axis=1) #batch, dim
return context, attention_values
def verify_compatible_shapes(img1, img2):
"""Checks if two image tensors are compatible for metric computation.
This function checks if two sets of images have ranks at least 3, and if the
last three dimensions match.
Args:
img1: The first images tensor.
img2: The second images tensor.
Returns:
A tuple of the first tensor shape, the second tensor shape, and a list of
tf.Assert() implementing the checks.
Raises:
ValueError: when static shape check fails.
"""
shape1 = img1.get_shape().with_rank_at_least(3)
shape2 = img2.get_shape().with_rank_at_least(3)
shape1[-3:].assert_is_compatible_with(shape2[-3:])
if shape1.ndims is not None and shape2.ndims is not None:
for dim1, dim2 in zip(reversed(shape1[:-3]), reversed(shape2[:-3])):
if not (dim1 == 1 or dim2 == 1 or dim1.is_compatible_with(dim2)):
raise ValueError(
'Two images are not compatible: %s and %s' % (shape1, shape2))
# Now assign shape tensors.
shape1, shape2 = tf.shape_n([img1, img2])
checks = []
checks.append(tf.Assert(tf.greater_equal(tf.size(shape1), 3),
[shape1, shape2], summarize=10))
checks.append(tf.Assert(tf.reduce_all(tf.equal(shape1[-3:], shape2[-3:])),
[shape1, shape2], summarize=10))
return shape1, shape2, checks
def Scipy_Keras_Wrapper(model, loss_func, X_i, u_i, X_b, u_b, X_f):
#Borrowed from Py Chao https://pychao.com/2019/11/02/optimize-tensorflow-keras-models-with-l-bfgs-from-tensorflow-probability/
# obtain the shapes of all trainable parameters in the model
shapes = tf.shape_n(model.trainable_variables)
n_tensors = len(shapes)
# we'll use tf.dynamic_stitch and tf.dynamic_partition later, so we need to
# prepare required information first
count = 0
idx = [] # stitch indices
part = [] # partition indices
loss = []
for i, shape in enumerate(shapes):
n = np.product(shape)
idx.append(
tf.reshape(tf.range(count, count + n, dtype=tf.int32), shape))
part.extend([i] * n)
count += n
part = tf.constant(part)
def assign_new_model_parameters(params_1d):
"""A function updating the model's parameters with a 1D tf.Tensor.
Args:
params_1d [in]: a 1D tf.Tensor representing the model's trainable parameters.
"""
params = tf.dynamic_partition(params_1d, part, n_tensors)
for i, (shape, param) in enumerate(zip(shapes, params)):
model.trainable_variables[i].assign(tf.reshape(param, shape))
def val_and_grads_1d(params_1d):
"""A function that can be used by tfp.optimizer.lbfgs_minimize.
This function is created by function_factory.
Args:
params_1d [in]: a 1D tf.Tensor.
Returns:
A scalar loss and the gradients w.r.t. the `params_1d`.
"""
start_time = time.time()
# use GradientTape so that we can calculate the gradient of loss w.r.t. parameters
with tf.GradientTape() as tape:
# update the parameters in the model
assign_new_model_parameters(params_1d)
# calculate the loss
loss_value = loss_func(X_i, u_i, X_b, u_b, X_f)
# calculate gradients and convert to 1D tf.Tensor
grads = tape.gradient(loss_value, model.trainable_variables)
grads = tf.dynamic_stitch(idx, grads)
# print out iteration & loss
val_and_grads_1d.iter.assign_add(1)
print('QN. It: %d, Loss: %.3e, Time: %.2f' %
(val_and_grads_1d.iter, loss_value,
np.round(time.time() - start_time, 3)))
# tf.print("Iter:", val_and_grads_1d.iter, "Loss:", np.round(loss_value, 5), "Time:", np.round(time.time() - start_time, 2))
return loss_value.numpy(), grads.numpy()
# store these information as members so we can use them outside the scope
val_and_grads_1d.iter = tf.Variable(0)
val_and_grads_1d.idx = idx
val_and_grads_1d.part = part
val_and_grads_1d.shapes = shapes
val_and_grads_1d.assign_new_model_parameters = assign_new_model_parameters
val_and_grads_1d.loss = loss
return val_and_grads_1d
def _ssim_per_channel(img1, img2, max_val=1.0):
"""
Computes SSIM index between img1 and img2 per color channel.
This function matches the standard SSIM implementation from:
Wang, Z., Bovik, A. C., Sheikh, H. R., & Simoncelli, E. P. (2004). Image
quality assessment: from error visibility to structural similarity. IEEE
transactions on image processing.
Details:
- 11x11 Gaussian filter of width 1.5 is used.
- k1 = 0.01, k2 = 0.03 as in the original paper.
Args:
img1: First image batch.
img2: Second image batch.
max_val: The dynamic range of the images (i.e., the difference between the
maximum the and minimum allowed values).
Returns:
A pair of tensors containing and channel-wise SSIM and contrast-structure
values. The shape is [..., channels].
"""
def _fspecial_gauss(size, sigma):
"""
Function to mimic the 'fspecial' gaussian MATLAB function.
"""
size = tf.convert_to_tensor(size, 'int32')
sigma = tf.convert_to_tensor(sigma)
coords = tf.cast(tf.range(size), sigma.dtype)
coords -= tf.cast(size - 1, sigma.dtype) / 2.0
var_g = tf.square(coords)
var_g *= -0.5 / tf.square(sigma)
var_g = tf.reshape(var_g, shape=[1, -1]) + tf.reshape(var_g, shape=[-1, 1])
var_g = tf.reshape(var_g, shape=[1, -1]) # For tf.nn.softmax().
var_g = tf.nn.softmax(var_g)
return tf.reshape(var_g, shape=[size, size, 1, 1])
def _ssim_helper(var_x, var_y, max_val, kernel, compensation=1.0):
"""
Helper function for computing SSIM.
SSIM estimates covariances with weighted sums. The default parameters
use a biased estimate of the covariance:
Suppose `reducer` is a weighted sum, then the mean estimators are
mu_x = sum_i w_i x_i,
mu_y = sum_i w_i y_i,
where w_i's are the weighted-sum weights, and covariance estimator is
cov_{xy} = sum_i w_i (x_i - mu_x) (y_i - mu_y)
with assumption sum_i w_i = 1. This covariance estimator is biased, since
E[cov_{xy}] = (1 - sum_i w_i ^ 2) Cov(X, Y).
For SSIM measure with unbiased covariance estimators, pass as `compensation`
argument (1 - sum_i w_i ^ 2).
Arguments:
x: First set of images.
y: Second set of images.
reducer: Function that computes 'local' averages from set of images.
For non-covolutional version, this is usually tf.reduce_mean(x, [1, 2]),
and for convolutional version, this is usually tf.nn.avg_pool or
tf.nn.conv2d with weighted-sum kernel.
max_val: The dynamic range (i.e., the difference between the maximum
possible allowed value and the minimum allowed value).
compensation: Compensation factor. See above.
Returns:
A pair containing the luminance measure, and the contrast-structure measure.
"""
def reducer(var_x, kernel):
shape = tf.shape(var_x)
var_x = tf.reshape(var_x, shape=tf.concat([[-1], shape[-3:]], 0))
var_y = tf.nn.depthwise_conv2d(var_x, kernel, strides=[1, 1, 1, 1], padding='VALID')
return tf.reshape(var_y, tf.concat([shape[:-3], tf.shape(var_y)[1:]], 0))
_ssim_k1 = 0.01
_ssim_k2 = 0.03
c_1 = (_ssim_k1 * max_val) ** 2
c_2 = (_ssim_k2 * max_val) ** 2
# SSIM luminance measure is
# (2 * mu_x * mu_y + c_1) / (mu_x ** 2 + mu_y ** 2 + c_1).
mean0 = reducer(var_x, kernel)
mean1 = reducer(var_y, kernel)
num0 = mean0 * mean1 * 2.0
den0 = tf.square(mean0) + tf.square(mean1)
luminance = (num0 + c_1) / (den0 + c_1)
# SSIM contrast-structure measure is
# (2 * cov_{xy} + c_2) / (cov_{xx} + cov_{yy} + c_2).
# Note that `reducer` is a weighted sum with weight w_k, \sum_i w_i = 1, then
# cov_{xy} = \sum_i w_i (x_i - \mu_x) (y_i - \mu_y)
# = \sum_i w_i x_i y_i - (\sum_i w_i x_i) (\sum_j w_j y_j).
num1 = reducer(var_x * var_y, kernel) * 2.0
den1 = reducer(tf.square(var_x) + tf.square(var_y), kernel)
c_2 *= compensation
c_s = (num1 - num0 + c_2) / (den1 - den0 + c_2)
# SSIM score is the product of the luminance and contrast-structure measures.
return luminance, c_s
filter_size = tf.constant(9, dtype='int32') # changed from 11 to 9 due
filter_sigma = tf.constant(1.5, dtype=img1.dtype)
shape1, shape2 = tf.shape_n([img1, img2])
checks = [tf.Assert(tf.reduce_all(tf.greater_equal(shape1[-3:-1], filter_size)),
[shape1, filter_size], summarize=8),
tf.Assert(tf.reduce_all(tf.greater_equal(shape2[-3:-1], filter_size)),
[shape2, filter_size], summarize=8)]
# Enforce the check to run before computation.
with tf.control_dependencies(checks):
img1 = tf.identity(img1)
# TODO(sjhwang): Try to cache kernels and compensation factor.
kernel = _fspecial_gauss(filter_size, filter_sigma)
kernel = tf.tile(kernel, multiples=[1, 1, shape1[-1], 1])
# The correct compensation factor is `1.0 - tf.reduce_sum(tf.square(kernel))`,
# but to match MATLAB implementation of MS-SSIM, we use 1.0 instead.
compensation = 1.0
# TODO(sjhwang): Try FFT.
# TODO(sjhwang): Gaussian kernel is separable in space. Consider applying
# 1-by-n and n-by-1 Gaussain filters instead of an n-by-n filter.
luminance, c_s = _ssim_helper(img1, img2, max_val, kernel, compensation)
# Average over the second and the third from the last: height, width.
axes = tf.constant([-3, -2], dtype='int32')
ssim_val = tf.reduce_mean(luminance * c_s, axes)
c_s = tf.reduce_mean(c_s, axes)
return ssim_val, c_s
def function_factory(model, residual, loss_bc, loss_ic, x, y, t, x_train_bc,
y_train_bc, t_train_bc, bc_train, x_train_ic, y_train_ic,
t_train_ic, ic_train, alpha, beta, D):
"""A factory to create a function required by tfp.optimizer.lbfgs_minimize.
Args:
model [in]: an instance of `tf.keras.Model` or its subclasses.
loss [in]: a function with signature loss_value = loss(pred_y, true_y).
train_x [in]: the input part of training data.
train_y [in]: the output part of training data.
Returns:
A function that has a signature of:
loss_value, gradients = f(model_parameters).
"""
# 计算model中可训练参数的形状
shapes = tf.shape_n(model.trainable_variables)
n_tensors = len(shapes)
# 随后将使用tf.dynamic_switch和tf.dynamic_partition来更改形状，下面是准备工作
count = 0
idx = [] # stitch indices
part = [] # partition indices
for i, shape in enumerate(shapes):
n = np.product(shape)
idx.append(
tf.reshape(tf.range(count, count + n, dtype=tf.int32), shape))
part.extend([i] * n)
count += n
part = tf.constant(part)
@tf.function
def assign_new_model_parameters(params_1d):
"""A function updating the model's parameters with a 1D tf.Tensor.
Args:
params_1d [in]: a 1D tf.Tensor representing the model's trainable parameters.
"""
params = tf.dynamic_partition(params_1d, part, n_tensors)
for i, (shape, param) in enumerate(zip(shapes, params)):
model.trainable_variables[i].assign(tf.reshape(param, shape))
# now create a function that will be returned by this factory
@tf.function
def f(params_1d):
"""A function that can be used by tfp.optimizer.lbfgs_minimize.
This function is created by function_factory.
Args:
params_1d [in]: a 1D tf.Tensor.
Returns:
A scalar loss and the gradients the `params_1d`.
"""
# use GradientTape so that we can calculate the gradient of loss
with tf.GradientTape(persistent=True) as tape1:
# update the parameters in the model
assign_new_model_parameters(params_1d)
with tf.GradientTape(persistent=True) as tape2:
with tf.GradientTape(persistent=True) as tape3:
Psi = model(tf.concat([x, y, t], 1))
Psi_x, Psi_y, Psi_t = tape3.gradient(Psi, [x, y, t])
Psi_xx = tape2.gradient(Psi_x, x)
Psi_yy = tape2.gradient(Psi_y, y)
loss_e = residual(Psi_xx=Psi_xx,
Psi_yy=Psi_yy,
Psi_x=Psi_x,
Psi_y=Psi_y,
Psi_t=Psi_t,
x=x,
y=y,
t=t,
D=D)
loss_b = loss_bc(model=model,
x=x_train_bc,
y=y_train_bc,
t=t_train_bc,
bc_train=bc_train)
loss_i = loss_ic(model=model,
x=x_train_ic,
y=y_train_ic,
t=t_train_ic,
ic_train=ic_train)
# calculate the loss
loss_value = loss_e + alpha * loss_b + beta * loss_i
# calculate gradients and convert to 1D tf.Tensor
grads = tape1.gradient(loss_value, model.trainable_variables)
grads = tf.dynamic_stitch(idx, grads)
# print out iteration & loss
f.iter.assign_add(1)
if f.iter % 2000 == 0:
tf.print("Iter:", f.iter, "loss:", loss_value)
# store loss value so we can retrieve later
tf.py_function(f.history.append,
inp=[[loss_value, loss_e, loss_b, loss_i]],
Tout=[])
return loss_value, grads
# store these information as members so we can use them outside the scope
f.iter = tf.Variable(0)
f.idx = idx
f.part = part
f.shapes = shapes
f.assign_new_model_parameters = assign_new_model_parameters
f.history = []
return f
def ms_ssim(img1, img2, max_val=1.0, power_factors=(0.0517, 0.3295, 0.3462, 0.2726)):
"""
Computes the MS-SSIM between img1 and img2.
This function assumes that `img1` and `img2` are image batches, i.e. the last
three dimensions are [height, width, channels].
Note: The true SSIM is only defined on grayscale. This function does not
perform any colorspace transform. (If input is already YUV, then it will
compute YUV SSIM average.)
Original paper: Wang, Zhou, Eero P. Simoncelli, and Alan C. Bovik. "Multiscale
structural similarity for image quality assessment." Signals, Systems and
Computers, 2004.
Arguments:
img1: First image batch.
img2: Second image batch. Must have the same rank as img1.
max_val: The dynamic range of the images (i.e., the difference between the
maximum the and minimum allowed values).
power_factors: Iterable of weights for each of the scales. The number of
scales used is the length of the list. Index 0 is the unscaled
resolution's weight and each increasing scale corresponds to the image
being downsampled by 2. Defaults to (0.0448, 0.2856, 0.3001, 0.2363,
0.1333), which are the values obtained in the original paper.
Returns:
A tensor containing an MS-SSIM value for each image in batch. The values
are in range [0, 1]. Returns a tensor with shape:
broadcast(img1.shape[:-3], img2.shape[:-3]).
"""
def _verify_compatible_image_shapes(img1, img2):
"""
Checks if two image tensors are compatible for applying SSIM or PSNR.
This function checks if two sets of images have ranks at least 3, and if the
last three dimensions match.
Args:
img1: Tensor containing the first image batch.
img2: Tensor containing the second image batch.
Returns:
A tuple containing: the first tensor shape, the second tensor shape, and a
list of control_flow_ops.Assert() ops implementing the checks.
Raises:
ValueError: When static shape check fails.
"""
shape1 = img1.get_shape().with_rank_at_least(3)
shape2 = img2.get_shape().with_rank_at_least(3)
shape1[-3:].assert_is_compatible_with(shape2[-3:])
if shape1.ndims is not None and shape2.ndims is not None:
for dim1, dim2 in zip(reversed(shape1[:-3]), reversed(shape2[:-3])):
if not (dim1 == 1 or dim2 == 1 or dim1.is_compatible_with(dim2)):
raise ValueError('Two images are not compatible: %s and %s' % (shape1, shape2))
# Now assign shape tensors.
shape1, shape2 = tf.shape_n([img1, img2])
# TODO(sjhwang): Check if shape1[:-3] and shape2[:-3] are broadcastable.
checks = []
checks.append(tf.Assert(tf.greater_equal(tf.size(shape1), 3),
[shape1, shape2], summarize=10))
checks.append(tf.Assert(tf.reduce_all(tf.equal(shape1[-3:], shape2[-3:])),
[shape1, shape2], summarize=10))
return shape1, shape2, checks
def _ssim_per_channel(img1, img2, max_val=1.0):
"""
Computes SSIM index between img1 and img2 per color channel.
This function matches the standard SSIM implementation from:
Wang, Z., Bovik, A. C., Sheikh, H. R., & Simoncelli, E. P. (2004). Image
quality assessment: from error visibility to structural similarity. IEEE
transactions on image processing.
Details:
- 11x11 Gaussian filter of width 1.5 is used.
- k1 = 0.01, k2 = 0.03 as in the original paper.
Args:
img1: First image batch.
img2: Second image batch.
max_val: The dynamic range of the images (i.e., the difference between the
maximum the and minimum allowed values).
Returns:
A pair of tensors containing and channel-wise SSIM and contrast-structure
values. The shape is [..., channels].
"""
def _fspecial_gauss(size, sigma):
"""
Function to mimic the 'fspecial' gaussian MATLAB function.
"""
size = tf.convert_to_tensor(size, 'int32')
sigma = tf.convert_to_tensor(sigma)
coords = tf.cast(tf.range(size), sigma.dtype)
coords -= tf.cast(size - 1, sigma.dtype) / 2.0
var_g = tf.square(coords)
var_g *= -0.5 / tf.square(sigma)
var_g = tf.reshape(var_g, shape=[1, -1]) + tf.reshape(var_g, shape=[-1, 1])
var_g = tf.reshape(var_g, shape=[1, -1]) # For tf.nn.softmax().
var_g = tf.nn.softmax(var_g)
return tf.reshape(var_g, shape=[size, size, 1, 1])
def _ssim_helper(var_x, var_y, max_val, kernel, compensation=1.0):
"""
Helper function for computing SSIM.
SSIM estimates covariances with weighted sums. The default parameters
use a biased estimate of the covariance:
Suppose `reducer` is a weighted sum, then the mean estimators are
mu_x = sum_i w_i x_i,
mu_y = sum_i w_i y_i,
where w_i's are the weighted-sum weights, and covariance estimator is
cov_{xy} = sum_i w_i (x_i - mu_x) (y_i - mu_y)
with assumption sum_i w_i = 1. This covariance estimator is biased, since
E[cov_{xy}] = (1 - sum_i w_i ^ 2) Cov(X, Y).
For SSIM measure with unbiased covariance estimators, pass as `compensation`
argument (1 - sum_i w_i ^ 2).
Arguments:
x: First set of images.
y: Second set of images.
reducer: Function that computes 'local' averages from set of images.
For non-covolutional version, this is usually tf.reduce_mean(x, [1, 2]),
and for convolutional version, this is usually tf.nn.avg_pool or
tf.nn.conv2d with weighted-sum kernel.
max_val: The dynamic range (i.e., the difference between the maximum
possible allowed value and the minimum allowed value).
compensation: Compensation factor. See above.
Returns:
A pair containing the luminance measure, and the contrast-structure measure.
"""
def reducer(var_x, kernel):
shape = tf.shape(var_x)
var_x = tf.reshape(var_x, shape=tf.concat([[-1], shape[-3:]], 0))
var_y = tf.nn.depthwise_conv2d(var_x, kernel, strides=[1, 1, 1, 1], padding='VALID')
return tf.reshape(var_y, tf.concat([shape[:-3], tf.shape(var_y)[1:]], 0))
_ssim_k1 = 0.01
_ssim_k2 = 0.03
c_1 = (_ssim_k1 * max_val) ** 2
c_2 = (_ssim_k2 * max_val) ** 2
# SSIM luminance measure is
# (2 * mu_x * mu_y + c_1) / (mu_x ** 2 + mu_y ** 2 + c_1).
mean0 = reducer(var_x, kernel)
mean1 = reducer(var_y, kernel)
num0 = mean0 * mean1 * 2.0
den0 = tf.square(mean0) + tf.square(mean1)
luminance = (num0 + c_1) / (den0 + c_1)
# SSIM contrast-structure measure is
# (2 * cov_{xy} + c_2) / (cov_{xx} + cov_{yy} + c_2).
# Note that `reducer` is a weighted sum with weight w_k, \sum_i w_i = 1, then
# cov_{xy} = \sum_i w_i (x_i - \mu_x) (y_i - \mu_y)
# = \sum_i w_i x_i y_i - (\sum_i w_i x_i) (\sum_j w_j y_j).
num1 = reducer(var_x * var_y, kernel) * 2.0
den1 = reducer(tf.square(var_x) + tf.square(var_y), kernel)
c_2 *= compensation
c_s = (num1 - num0 + c_2) / (den1 - den0 + c_2)
# SSIM score is the product of the luminance and contrast-structure measures.
return luminance, c_s
filter_size = tf.constant(9, dtype='int32') # changed from 11 to 9 due
filter_sigma = tf.constant(1.5, dtype=img1.dtype)
shape1, shape2 = tf.shape_n([img1, img2])
checks = [tf.Assert(tf.reduce_all(tf.greater_equal(shape1[-3:-1], filter_size)),
[shape1, filter_size], summarize=8),
tf.Assert(tf.reduce_all(tf.greater_equal(shape2[-3:-1], filter_size)),
[shape2, filter_size], summarize=8)]
# Enforce the check to run before computation.
with tf.control_dependencies(checks):
img1 = tf.identity(img1)
# TODO(sjhwang): Try to cache kernels and compensation factor.
kernel = _fspecial_gauss(filter_size, filter_sigma)
kernel = tf.tile(kernel, multiples=[1, 1, shape1[-1], 1])
# The correct compensation factor is `1.0 - tf.reduce_sum(tf.square(kernel))`,
# but to match MATLAB implementation of MS-SSIM, we use 1.0 instead.
compensation = 1.0
# TODO(sjhwang): Try FFT.
# TODO(sjhwang): Gaussian kernel is separable in space. Consider applying
# 1-by-n and n-by-1 Gaussain filters instead of an n-by-n filter.
luminance, c_s = _ssim_helper(img1, img2, max_val, kernel, compensation)
# Average over the second and the third from the last: height, width.
axes = tf.constant([-3, -2], dtype='int32')
ssim_val = tf.reduce_mean(luminance * c_s, axes)
c_s = tf.reduce_mean(c_s, axes)
return ssim_val, c_s
def do_pad(images, remainder):
padding = tf.expand_dims(remainder, -1)
padding = tf.pad(padding, [[1, 0], [1, 0]])
return [tf.pad(x, padding, mode='SYMMETRIC') for x in images]
# Shape checking.
shape1 = img1.get_shape().with_rank_at_least(3)
shape2 = img2.get_shape().with_rank_at_least(3)
shape1[-3:].merge_with(shape2[-3:])
with tf.name_scope(None, 'MS-SSIM', [img1, img2]):
shape1, shape2, checks = _verify_compatible_image_shapes(img1, img2)
with tf.control_dependencies(checks):
img1 = tf.identity(img1)
# Need to convert the images to float32. Scale max_val accordingly so that
# SSIM is computed correctly.
max_val = tf.cast(max_val, img1.dtype)
max_val = tf.image.convert_image_dtype(max_val, 'float32')
img1 = tf.image.convert_image_dtype(img1, 'float32')
img2 = tf.image.convert_image_dtype(img2, 'float32')
imgs = [img1, img2]
shapes = [shape1, shape2]
# img1 and img2 are assumed to be a (multi-dimensional) batch of
# 3-dimensional images (height, width, channels). `heads` contain the batch
# dimensions, and `tails` contain the image dimensions.
heads = [s[:-3] for s in shapes]
tails = [s[-3:] for s in shapes]
divisor = [1, 2, 2, 1]
divisor_tensor = tf.constant(divisor[1:], dtype='int32')
mc_s = []
for k in range(len(power_factors)):
with tf.name_scope(None, 'Scale%d' % k, imgs):
if k > 0:
# Avg pool takes rank 4 tensors. Flatten leading dimensions.
flat_imgs = [tf.reshape(x, tf.concat([[-1], t], 0))
for x, t in zip(imgs, tails)]
remainder = tails[0] % divisor_tensor
need_padding = tf.reduce_any(tf.not_equal(remainder, 0))
padded = tf.cond(need_padding,
lambda: do_pad(flat_imgs, remainder),
lambda: flat_imgs)
downscaled = [tf.nn.avg_pool(x, ksize=divisor, strides=divisor, padding='VALID')
for x in padded]
tails = [x[1:] for x in tf.shape_n(downscaled)]
imgs = [tf.reshape(x, tf.concat([h, t], 0))
for x, h, t in zip(downscaled, heads, tails)]
# Overwrite previous ssim value since we only need the last one.
ssim_per_channel, c_s = _ssim_per_channel(*imgs, max_val=max_val)
mc_s.append(tf.nn.relu(c_s))
# Remove the c_s score for the last scale. In the MS-SSIM calculation,
# we use the l(p) at the highest scale. l(p) * c_s(p) is ssim(p).
mc_s.pop() # Remove the c_s score for the last scale.
mc_s_and_ssim = tf.stack(mc_s + [tf.nn.relu(ssim_per_channel)], axis=-1)
# Take weighted geometric mean across the scale axis.
ms_ssim = tf.reduce_prod(tf.pow(mc_s_and_ssim, power_factors), [-1])
return tf.reduce_mean(ms_ssim, [-1]) # Avg over color channels.
def function_factory(model, loss, train_x, train_y, batch_size=10):
# obtain the shapes of all trainable parameters in the model
shapes = tf.shape_n(model.trainable_variables)
n_tensors = len(shapes)
# we'll use tf.dynamic_stitch and tf.dynamic_partition later, so we need to
# prepare required information first
count = 0
idx = [] # stitch indices
part = [] # partition indices
for i, shape in enumerate(shapes):
n = np.product(shape)
idx.append(
tf.reshape(tf.range(count, count + n, dtype=tf.int32), shape))
part.extend([i] * n)
count += n
part = tf.constant(part)
@tf.function
def assign_new_model_parameters(params_1d):
params = tf.dynamic_partition(params_1d, part, n_tensors)
for i, (shape, param) in enumerate(zip(shapes, params)):
model.trainable_variables[i].assign(tf.reshape(param, shape))
# now create a function that will be returned by this factory
@tf.function
def f(params_1d):
N = len(x_train)
# use GradientTape so that we can calculate the gradient of loss w.r.t. parameters
with tf.GradientTape() as tape:
# update the parameters in the model
assign_new_model_parameters(params_1d)
if batch_size < N:
perm = np.random.permutation(N)
loss_value = 0
for i in range(0, N, batch_size):
# calculate the loss
loss_value += loss(
model(train_x[perm[i:i + batch_size]], training=True),
train_y[perm[i:i + batch_size]])
else:
# calculate the loss
loss_value = loss(model(train_x, training=True), train_y)
# calculate gradients and convert to 1D tf.Tensor
grads = tape.gradient(loss_value, model.trainable_variables)
grads = tf.dynamic_stitch(idx, grads)
# print out iteration & loss
f.iter.assign_add(1)
tf.print("Iter:", f.iter, "loss:", loss_value)
# store loss value so we can retrieve later
tf.py_function(f.history.append, inp=[loss_value], Tout=[])
return loss_value, grads
# store these information as members so we can use them outside the scope
f.iter = tf.Variable(0)
f.idx = idx
f.part = part
f.shapes = shapes
f.assign_new_model_parameters = assign_new_model_parameters
f.history = []
return f
def ms_ssim_calc(img1,
img2,
max_val=1.0,
power_factors=(0.0517, 0.3295, 0.3462, 0.2726)):
"""
Computes the MS-SSIM between img1 and img2.
This function assumes that `img1` and `img2` are image batches, i.e. the last
three dimensions are [height, width, channels].
Note: The true SSIM is only defined on grayscale. This function does not
perform any colorspace transform. (If input is already YUV, then it will
compute YUV SSIM average.)
Original paper: Wang, Zhou, Eero P. Simoncelli, and Alan C. Bovik. "Multiscale
structural similarity for image quality assessment." Signals, Systems and
Computers, 2004.
Arguments:
img1: First image batch.
img2: Second image batch. Must have the same rank as img1.
max_val: The dynamic range of the images (i.e., the difference between the
maximum the and minimum allowed values).
power_factors: Iterable of weights for each of the scales. The number of
scales used is the length of the list. Index 0 is the unscaled
resolution's weight and each increasing scale corresponds to the image
being downsampled by 2. Defaults to (0.0448, 0.2856, 0.3001, 0.2363,
0.1333), which are the values obtained in the original paper.
Returns:
A tensor containing an MS-SSIM value for each image in batch. The values
are in range [0, 1]. Returns a tensor with shape:
broadcast(img1.shape[:-3], img2.shape[:-3]).
"""
def _verify_compatible_image_shapes(img1, img2):
"""
Checks if two image tensors are compatible for applying SSIM or PSNR.
This function checks if two sets of images have ranks at least 3, and if the
last three dimensions match.
Args:
img1: Tensor containing the first image batch.
img2: Tensor containing the second image batch.
Returns:
A tuple containing: the first tensor shape, the second tensor shape, and a
list of control_flow_ops.Assert() ops implementing the checks.
Raises:
ValueError: When static shape check fails.
"""
shape1 = img1.get_shape().with_rank_at_least(3)
shape2 = img2.get_shape().with_rank_at_least(3)
shape1[-3:].assert_is_compatible_with(shape2[-3:])
if shape1.ndims is not None and shape2.ndims is not None:
for dim1, dim2 in zip(reversed(shape1[:-3]),
reversed(shape2[:-3])):
if not (dim1 == 1 or dim2 == 1
or dim1.is_compatible_with(dim2)):
raise ValueError(
'Two images are not compatible: %s and %s' %
(shape1, shape2))
# Now assign shape tensors.
shape1, shape2 = tf.shape_n([img1, img2])
# TODO(sjhwang): Check if shape1[:-3] and shape2[:-3] are broadcastable.
checks = []
checks.append(
tf.Assert(tf.greater_equal(tf.size(shape1), 3), [shape1, shape2],
summarize=10))
checks.append(
tf.Assert(tf.reduce_all(tf.equal(shape1[-3:], shape2[-3:])),
[shape1, shape2],
summarize=10))
return shape1, shape2, checks
def _ssim_per_channel(img1, img2, max_val=1.0):
"""
Computes SSIM index between img1 and img2 per color channel.
This function matches the standard SSIM implementation from:
Wang, Z., Bovik, A. C., Sheikh, H. R., & Simoncelli, E. P. (2004). Image
quality assessment: from error visibility to structural similarity. IEEE
transactions on image processing.
Details:
- 11x11 Gaussian filter of width 1.5 is used.
- k1 = 0.01, k2 = 0.03 as in the original paper.
Args:
img1: First image batch.
img2: Second image batch.
max_val: The dynamic range of the images (i.e., the difference between the
maximum the and minimum allowed values).
Returns:
A pair of tensors containing and channel-wise SSIM and contrast-structure
values. The shape is [..., channels].
"""
def _fspecial_gauss(size, sigma):
""" Function to mimic the 'fspecial' gaussian MATLAB function. """
size = tf.convert_to_tensor(size, 'int32')
sigma = tf.convert_to_tensor(sigma)
coords = tf.cast(tf.range(size), sigma.dtype)
coords -= tf.cast(size - 1, sigma.dtype) / 2.0
gauss = tf.square(coords)
gauss *= -0.5 / tf.square(sigma)
gauss = tf.reshape(gauss, shape=[1, -1]) + tf.reshape(
gauss, shape=[-1, 1])
gauss = tf.reshape(gauss, shape=[1, -1]) # For tf.nn.softmax().
gauss = tf.nn.softmax(gauss)
return tf.reshape(gauss, shape=[size, size, 1, 1])
def _ssim_helper(img1, img2, max_val, kernel, compensation=1.):
"""
Helper function for computing SSIM.
SSIM estimates covariances with weighted sums. The default parameters
use a biased estimate of the covariance:
Suppose `reducer` is a weighted sum, then the mean estimators are
mu_x = sum_i w_i x_i,
mu_y = sum_i w_i y_i,
where w_i's are the weighted-sum weights, and covariance estimator is
cov_{xy} = sum_i w_i (x_i - mu_x) (y_i - mu_y)
with assumption sum_i w_i = 1. This covariance estimator is biased, since
E[cov_{xy}] = (1 - sum_i w_i ^ 2) Cov(X, Y).
For SSIM measure with unbiased covariance estimators, pass as `compensation`
argument (1 - sum_i w_i ^ 2).
Arguments:
img1: First set of images.
img2: Second set of images.
reducer: Function that computes 'local' averages from set of images.
For non-covolutional version, this is usually tf.reduce_mean(img1, [1, 2]),
and for convolutional version, this is usually tf.nn.avg_pool or
tf.nn.conv2d with weighted-sum kernel.
max_val: The dynamic range (i.e., the difference between the maximum
possible allowed value and the minimum allowed value).
compensation: Compensation factor. See above.
Returns:
A pair containing the luminance measure, and the contrast-structure measure.
"""
def reducer(img1, kernel):
shape = tf.shape(img1)
img1 = tf.reshape(img1, shape=tf.concat([[-1], shape[-3:]], 0))
img2 = tf.nn.depthwise_conv2d(img1,
kernel,
strides=[1, 1, 1, 1],
padding='VALID')
return tf.reshape(
img2,
tf.concat([shape[:-3], tf.shape(img2)[1:]], 0))
c_one = (0.01 * max_val)**2
c_two = ((0.03 * max_val))**2 * compensation
# SSIM luminance measure is
# (2 * mu_x * mu_y + c_one) / (mu_x ** 2 + mu_y ** 2 + c_one).
mean0 = reducer(img1, kernel)
mean1 = reducer(img2, kernel)
num0 = mean0 * mean1 * 2.
den0 = tf.square(mean0) + tf.square(mean1)
luminance = (num0 + c_one) / (den0 + c_one)
# SSIM contrast-structure measure is
# (2 * cov_{xy} + c_two) / (cov_{xx} + cov_{yy} + c_two).
# Note that `reducer` is a weighted sum with weight w_k, \sum_i w_i = 1, then
# cov_{xy} = \sum_i w_i (x_i - \mu_x) (y_i - \mu_y)
# = \sum_i w_i x_i y_i - (\sum_i w_i x_i) (\sum_j w_j y_j).
num1 = reducer(img1 * img2, kernel) * 2.0
den1 = reducer(tf.square(img1) + tf.square(img2), kernel)
c_s = (num1 - num0 + c_two) / (den1 - den0 + c_two)
# SSIM score is the product of the luminance and contrast-structure measures.
return luminance, c_s
filter_size = tf.constant(9, dtype='int32') # changed from 11 to 9 due
filter_sigma = tf.constant(1.5, dtype=img1.dtype)
shape1, shape2 = tf.shape_n([img1, img2])
checks = [
tf.Assert(tf.reduce_all(
tf.greater_equal(shape1[-3:-1], filter_size)),
[shape1, filter_size],
summarize=8),
tf.Assert(tf.reduce_all(
tf.greater_equal(shape2[-3:-1], filter_size)),
[shape2, filter_size],
summarize=8)
]
# Enforce the check to run before computation.
with tf.control_dependencies(checks):
img1 = tf.identity(img1)
# TODO(sjhwang): Try to cache kernels and compensation factor.
kernel = _fspecial_gauss(filter_size, filter_sigma)
kernel = tf.tile(kernel, multiples=[1, 1, shape1[-1], 1])
# The correct compensation factor is `1.0 - tf.reduce_sum(tf.square(kernel))`,
# but to match MATLAB implementation of MS-SSIM, we use 1.0 instead.
compensation = 1.
# TODO(sjhwang): Try FFT.
# TODO(sjhwang): Gaussian kernel is separable in space. Consider applying
# 1-by-n and n-by-1 Gaussain filters instead of an n-by-n filter.
luminance, c_s = _ssim_helper(img1, img2, max_val, kernel,
compensation)
# Average over the second and the third from the last: height, width.
axes = tf.constant([-3, -2], dtype='int32')
ssim_val = tf.reduce_mean(luminance * c_s, axes)
c_s = tf.reduce_mean(c_s, axes)
return ssim_val, c_s
def do_pad(images, remainder):
padding = tf.expand_dims(remainder, -1)
padding = tf.pad(padding, [[1, 0], [1, 0]])
return [tf.pad(x, padding, mode='SYMMETRIC') for x in images]
# Shape checking.
shape1 = img1.get_shape().with_rank_at_least(3)
shape2 = img2.get_shape().with_rank_at_least(3)
shape1[-3:].merge_with(shape2[-3:])
with tf.name_scope(None, 'MS-SSIM', [img1, img2]):
shape1, shape2, checks = _verify_compatible_image_shapes(img1, img2)
with tf.control_dependencies(checks):
img1 = tf.identity(img1)
# Need to convert the images to float32. Scale max_val accordingly so that
# SSIM is computed correctly.
max_val = tf.cast(max_val, img1.dtype)
max_val = tf.image.convert_image_dtype(max_val, 'float32')
img1 = tf.image.convert_image_dtype(img1, 'float32')
img2 = tf.image.convert_image_dtype(img2, 'float32')
imgs = [img1, img2]
shapes = [shape1, shape2]
# img1 and img2 are assumed to be a (multi-dimensional) batch of
# 3-dimensional images (height, width, channels). `heads` contain the batch
# dimensions, and `tails` contain the image dimensions.
heads = [s[:-3] for s in shapes]
tails = [s[-3:] for s in shapes]
divisor = [1, 2, 2, 1]
divisor_tensor = tf.constant(divisor[1:], dtype='int32')
mc_s = []
for k in range(len(power_factors)):
with tf.name_scope(None, 'Scale%d' % k, imgs):
if k > 0:
# Avg pool takes rank 4 tensors. Flatten leading dimensions.
zipped = zip(imgs, tails)
flat_imgs = [
tf.reshape(x, tf.concat([[-1], t], 0)) for x, t in zipped
]
remainder = tails[0] % divisor_tensor
need_padding = tf.reduce_any(tf.not_equal(remainder, 0))
padded = tf.cond(need_padding,
lambda: do_pad(flat_imgs, remainder),
lambda: flat_imgs)
downscaled = [
tf.nn.avg_pool(x,
ksize=divisor,
strides=divisor,
padding='VALID') for x in padded
]
tails = [x[1:] for x in tf.shape_n(downscaled)]
zipper = zip(downscaled, heads, tails)
imgs = [
tf.reshape(x, tf.concat([h, t], 0)) for x, h, t in zipper
]
# Overwrite previous ssim value since we only need the last one.
ssim_per_channel, c_s = _ssim_per_channel(*imgs, max_val=max_val)
mc_s.append(tf.nn.relu(c_s))
# Remove the c_s score for the last scale. In the MS-SSIM calculation,
# we use the l(p) at the highest scale. l(p) * c_s(p) is ssim(p).
mc_s.pop() # Remove the c_s score for the last scale.
mcs_and_ssim = tf.stack(mc_s + [tf.nn.relu(ssim_per_channel)], axis=-1)
# Take weighted geometric mean across the scale axis.
ms_ssim = tf.reduce_prod(tf.pow(mcs_and_ssim, power_factors), [-1])
return tf.reduce_mean(ms_ssim, [-1]) # Avg over color channels.
def function_factory(model, loss_f, X, Y,params,kernel):
"""A factory to create a function required by tfp.optimizer.lbfgs_minimize.
Args:
model [in]: an instance of `tf.keras.Model` or its subclasses.
loss [in]: a function with signature loss_value = loss(pred_y, true_y).
train_x [in]: the input part of training data.
train_y [in]: the output part of training data.
Returns:
A function that has a signature of:
loss_value, gradients = f(model_parameters).
"""
# obtain the shapes of all trainable parameters in the model
shapes = tf.shape_n(model._opti_variables)
n_tensors = len(shapes)
# we'll use tf.dynamic_stitch and tf.dynamic_partition later, so we need to
# prepare required information first
count = 0
idx = [] # stitch indices
part = [] # partition indices
for i, shape in enumerate(shapes):
n = np.product(shape)
idx.append(tf.reshape(tf.range(count, count+n, dtype=tf.int32), shape))
part.extend([i]*n)
count += n
part = tf.constant(part)
def assign_new_model_parameters(params_1d):
"""A function updating the model's parameters with a 1D tf.Tensor.
Args:
params_1d [in]: a 1D tf.Tensor representing the model's trainable parameters.
"""
params = tf.dynamic_partition(params_1d, part, n_tensors)
for i, (shape, param) in enumerate(zip(shapes, params)):
model._opti_variables[i].assign(tf.cast(tf.reshape(param, shape), dtype=_precision))
# now create a function that will be returned by this factory
def f(params_1d):
"""A function that can be used by tfp.optimizer.lbfgs_minimize.
This function is created by function_factory.
Args:
params_1d [in]: a 1D tf.Tensor.
Returns:
A scalar loss and the gradients w.r.t. the `params_1d`.
"""
# use GradientTape so that we can calculate the gradient of loss w.r.t. parameters
with tf.GradientTape() as tape:
# update the parameters in the model
assign_new_model_parameters(params_1d)
# calculate the loss
loss_value = model(X,Y,kernel)[0][0]
# calculate gradients and convert to 1D tf.Tensor
grads = tape.gradient(loss_value, model.variables)
grads = tf.dynamic_stitch(idx, grads)
# print out iteration & loss
#f.iter.assign_add(1)
#tf.print("Iter:", f.iter, "loss:", loss_value)
# store loss value so we can retrieve later
tf.py_function(f.history.append, inp=[loss_value], Tout=[])
#return loss_value ,grads
return np.array(loss_value, order='F'),np.array(grads, order='F')
# store these information as members so we can use them outside the scope
f.iter = tf.convert_to_tensor(0)
f.idx = idx
f.part = part
f.shapes = shapes
f.assign_new_model_parameters = assign_new_model_parameters
f.history = []
return f
try:
import tensorflow as tf
except ImportError:
print("WARNING: Tensorflow not found, skipping test")
exit(0)
from dace.frontend.tensorflow import TFSession
myshape = [69, 96, 666]
num_inputs = 5
inpList = [tf.ones(myshape) for _ in range(num_inputs)]
sess_tf = tf.Session()
sess_dace = TFSession()
shapes_tf = sess_tf.run(tf.shape_n(inpList))
shapes_dace = sess_dace.run(tf.shape_n(inpList))
for dc, tf in zip(shapes_dace, shapes_tf):
try:
assert (dc == tf).all()
except (AssertionError):
print(dc)
print(tf)
def _MultiscaleSSIMHelper(img1,
img2,
filter_size,
filter_width,
power_factors,
max_val=255.0):
"""Computes the MS-SSIM between img1 and img2.
This function assumes that `img1` and `img2` are image batches, i.e. the last
three dimensions are [row, col, channels].
Arguments:
img1: First RGB image batch.
img2: Second RGB image batch. Must have the same rank as img1.
filter_size: An integer, the filter size of the Gaussian kernel used.
filter_width: A float, the filter width of the Gaussian kernel used.
power_factors: iterable of weightings for each of the scales. The number of
scales used is the length of the list. Index 0 is the unscaled
resolution's weighting and each increasing scale corresponds to the image
being downsampled by 2.
max_val: the dynamic range of the images (i.e., the difference between the
maximum the and minimum allowed values).
Returns:
A tensor containing batch-wise MS-SSIM measure. MS-SSIM has range [0, 1].
The shape is broadcast(img1.shape[:-3], img2.shape[:-3]).
"""
# Shape checking.
shape1 = img1.get_shape().with_rank_at_least(3)
shape2 = img2.get_shape().with_rank_at_least(3)
shape1[-3:].merge_with(shape2[-3:])
with tf.name_scope(None, 'MS-SSIM', [img1, img2]):
shape1, shape2, checks = VerifyCompatibleImageShapes(img1, img2)
with tf.control_dependencies(checks):
img1 = tf.identity(img1)
imgs = [img1, img2]
shapes = [shape1, shape2]
# img1 and img2 are assumed to be a (multi-dimensional) batch of
# 3-dimensional images (height, width, channels). `heads` contain the batch
# dimensions, and `tails` contain the image dimensions.
heads = [s[:-3] for s in shapes]
tails = [s[-3:] for s in shapes]
divisor = [1, 2, 2, 1]
divisor_tensor = tf.constant(divisor[1:], dtype=tf.int32)
def do_pad(images, remainder): # pylint: disable=invalid-name
padding = tf.expand_dims(remainder, -1)
padding = tf.pad(padding, [[1, 0], [1, 0]])
return [tf.pad(x, padding, mode='SYMMETRIC') for x in images]
mcs = []
for k in range(len(power_factors)):
with tf.name_scope(None, 'Scale%d' % k, imgs):
if k > 0:
# Avg pool takes rank 4 tensors. Flatten leading dimensions.
flat_imgs = [
tf.reshape(x, tf.concat([[-1], t], 0))
for x, t in zip(imgs, tails)
]
remainder = tails[0] % divisor_tensor
need_padding = tf.reduce_any(tf.not_equal(remainder, 0))
# pylint: disable=cell-var-from-loop
padded = tf.cond(need_padding,
lambda: do_pad(flat_imgs, remainder),
lambda: flat_imgs)
# pylint: enable=cell-var-from-loop
downscaled = [
tf.nn.avg_pool(x,
ksize=divisor,
strides=divisor,
padding='VALID') for x in padded
]
tails = [x[1:] for x in tf.shape_n(downscaled)]
imgs = [
tf.reshape(x, tf.concat([h, t], 0))
for x, h, t in zip(downscaled, heads, tails)
]
# Overwrite previous ssim value since we only need the last one.
ssim, cs = _SSIMIndexPerChannel(*imgs,
filter_size=filter_size,
filter_width=filter_width,
max_val=max_val)
mcs.append(tf.nn.relu(cs))
# Remove the cs score for the last scale. In the MS-SSIM calculation,
# we use the l(p) at the highest scale. l(p) * cs(p) is ssim(p).
mcs.pop() # Remove the cs score for the last scale.
mcs_and_ssim = tf.stack(mcs + [tf.nn.relu(ssim)], axis=-1)
# Take weighted geometric mean across the scale axis.
ms_ssim = tf.reduce_prod(tf.pow(mcs_and_ssim, power_factors), [-1])
ms_ssim = tf.reduce_mean(ms_ssim, [-1]) # Average over color channels.
return ms_ssim
def scipy_function_factory(model, *args):
"""A factory to create a function required by scipy.opimizer.
Based on the example from https://stackoverflow.com/questions/59029854/use-scipy-optimizer-with-tensorflow-2-0-for-neural-network-training
Args:
model [in]: an instance of `tf.keras.Model` or its subclasses.
*args: arguments to be passed to model.get_grads method
Returns:
A function that has a signature of:
loss_value, gradients = f(model_parameters).
"""
# obtain the shapes of all trainable parameters in the model
shapes = tf.shape_n(model.trainable_variables)
n_tensors = len(shapes)
# we'll use tf.dynamic_stitch and tf.dynamic_partition later, so we need to
# prepare required information first
count = 0
idx = [] # stitch indices
part = [] # partition indices
for i, shape in enumerate(shapes):
n = np.product(shape)
idx.append(
tf.reshape(tf.range(count, count + n, dtype=tf.int32), shape))
part.extend([i] * n)
count += n
part = tf.constant(part)
def assign_new_model_parameters(params_1d):
"""A function updating the model's parameters with a 1D tf.Tensor.
Args:
params_1d [in]: a 1D tf.Tensor representing the model's trainable parameters.
"""
params = tf.dynamic_partition(params_1d, part, n_tensors)
for i, (shape, param) in enumerate(zip(shapes, params)):
model.trainable_variables[i].assign(tf.reshape(param, shape))
# now create a function that will be returned by this factory
def f(params_1d):
"""A function that can be used by tfp.optimizer.lbfgs_minimize.
This function is created by function_factory.
Args:
params_1d [in]: a 1D tf.Tensor.
Returns:
A scalar loss and the gradients w.r.t. the `params_1d`.
"""
assign_new_model_parameters(params_1d)
loss_value, grads = model.get_grad(*args)
grads = tf.dynamic_stitch(idx, grads)
# print out iteration & loss
f.iter.assign_add(1)
if f.iter % model.print_epoch == 0:
tf.print("Iter:", f.iter, "loss:", loss_value)
return loss_value.numpy(), grads.numpy()
# store these information as members so we can use them outside the scope
f.iter = tf.Variable(0)
f.idx = idx
f.part = part
f.shapes = shapes
f.assign_new_model_parameters = assign_new_model_parameters
return f
def _function_wrapper(self, loss_func, model):
"""A factory to create a function required by tfp.optimizer.lbfgs_minimize.
Args:
loss_func: a function with signature loss_value = loss(model).
model: an instance of `tf.keras.Model` or its subclasses.
Returns:
A function that has a signature of:
loss_value, gradients = f(model_parameters).
"""
# obtain the shapes of all trainable parameters in the model
shapes = tf.shape_n(model.trainable_variables)
n_tensors = len(shapes)
# we'll use tf.dynamic_stitch and tf.dynamic_partition later, so we need to
# prepare required information first
count = 0
idx = [] # stitch indices
part = [] # partition indices
for i, shape in enumerate(shapes):
n = np.product(shape)
idx.append(
tf.reshape(tf.range(count, count + n, dtype=tf.int32), shape))
part.extend([i] * n)
count += n
part = tf.constant(part)
@tf.function
def assign_new_model_parameters(params_1d):
"""A function updating the model's parameters with a 1D tf.Tensor.
Args:
params_1d [in]: a 1D tf.Tensor representing the model's trainable parameters.
"""
params = tf.dynamic_partition(params_1d, part, n_tensors)
for i, (shape, param) in enumerate(zip(shapes, params)):
model.trainable_variables[i].assign(tf.reshape(param, shape))
# now create a function that will be returned by this factory
def f(params_1d):
"""A function that can be used by tfp.optimizer.lbfgs_minimize.
This function is created by function_factory.
Args:
params_1d [in]: a 1D tf.Tensor.
Returns:
A scalar loss and the gradients w.r.t. the `params_1d`.
"""
# use GradientTape so that we can calculate the gradient of loss w.r.t. parameters
with tf.GradientTape() as tape:
# update the parameters in the model
assign_new_model_parameters(params_1d)
# calculate the loss
loss_value = loss_func(model)
# calculate gradients and convert to 1D tf.Tensor
grads = tape.gradient(loss_value, model.trainable_variables)
grads = tf.dynamic_stitch(idx, grads)
# print out iteration & loss
f.iter.assign_add(1)
tf.print("Iter:", f.iter, "loss:", loss_value)
if self.callback_list is not None:
info_dict = {
'iter': f.iter,
'loss': loss_value,
'grad': grads,
}
for callback in self.callback_list:
callback(model, info_dict=info_dict)
return loss_value, grads
if self.trace_function:
f = tf.function(f)
# store these information as members so we can use them outside the scope
f.iter = tf.Variable(0, trainable=False)
f.idx = idx
f.part = part
f.shapes = shapes
f.assign_new_model_parameters = assign_new_model_parameters
return f
def _ssim_per_channel(img1, img2, max_val=1.0):
"""
Computes SSIM index between img1 and img2 per color channel.
This function matches the standard SSIM implementation from:
Wang, Z., Bovik, A. C., Sheikh, H. R., & Simoncelli, E. P. (2004). Image
quality assessment: from error visibility to structural similarity. IEEE
transactions on image processing.
Details:
- 11x11 Gaussian filter of width 1.5 is used.
- k1 = 0.01, k2 = 0.03 as in the original paper.
Args:
img1: First image batch.
img2: Second image batch.
max_val: The dynamic range of the images (i.e., the difference between the
maximum the and minimum allowed values).
Returns:
A pair of tensors containing and channel-wise SSIM and contrast-structure
values. The shape is [..., channels].
"""
def _fspecial_gauss(size, sigma):
""" Function to mimic the 'fspecial' gaussian MATLAB function. """
size = tf.convert_to_tensor(size, 'int32')
sigma = tf.convert_to_tensor(sigma)
coords = tf.cast(tf.range(size), sigma.dtype)
coords -= tf.cast(size - 1, sigma.dtype) / 2.0
gauss = tf.square(coords)
gauss *= -0.5 / tf.square(sigma)
gauss = tf.reshape(gauss, shape=[1, -1]) + tf.reshape(
gauss, shape=[-1, 1])
gauss = tf.reshape(gauss, shape=[1, -1]) # For tf.nn.softmax().
gauss = tf.nn.softmax(gauss)
return tf.reshape(gauss, shape=[size, size, 1, 1])
def _ssim_helper(img1, img2, max_val, kernel, compensation=1.):
"""
Helper function for computing SSIM.
SSIM estimates covariances with weighted sums. The default parameters
use a biased estimate of the covariance:
Suppose `reducer` is a weighted sum, then the mean estimators are
mu_x = sum_i w_i x_i,
mu_y = sum_i w_i y_i,
where w_i's are the weighted-sum weights, and covariance estimator is
cov_{xy} = sum_i w_i (x_i - mu_x) (y_i - mu_y)
with assumption sum_i w_i = 1. This covariance estimator is biased, since
E[cov_{xy}] = (1 - sum_i w_i ^ 2) Cov(X, Y).
For SSIM measure with unbiased covariance estimators, pass as `compensation`
argument (1 - sum_i w_i ^ 2).
Arguments:
img1: First set of images.
img2: Second set of images.
reducer: Function that computes 'local' averages from set of images.
For non-covolutional version, this is usually tf.reduce_mean(img1, [1, 2]),
and for convolutional version, this is usually tf.nn.avg_pool or
tf.nn.conv2d with weighted-sum kernel.
max_val: The dynamic range (i.e., the difference between the maximum
possible allowed value and the minimum allowed value).
compensation: Compensation factor. See above.
Returns:
A pair containing the luminance measure, and the contrast-structure measure.
"""
def reducer(img1, kernel):
shape = tf.shape(img1)
img1 = tf.reshape(img1, shape=tf.concat([[-1], shape[-3:]], 0))
img2 = tf.nn.depthwise_conv2d(img1,
kernel,
strides=[1, 1, 1, 1],
padding='VALID')
return tf.reshape(
img2,
tf.concat([shape[:-3], tf.shape(img2)[1:]], 0))
c_one = (0.01 * max_val)**2
c_two = ((0.03 * max_val))**2 * compensation
# SSIM luminance measure is
# (2 * mu_x * mu_y + c_one) / (mu_x ** 2 + mu_y ** 2 + c_one).
mean0 = reducer(img1, kernel)
mean1 = reducer(img2, kernel)
num0 = mean0 * mean1 * 2.
den0 = tf.square(mean0) + tf.square(mean1)
luminance = (num0 + c_one) / (den0 + c_one)
# SSIM contrast-structure measure is
# (2 * cov_{xy} + c_two) / (cov_{xx} + cov_{yy} + c_two).
# Note that `reducer` is a weighted sum with weight w_k, \sum_i w_i = 1, then
# cov_{xy} = \sum_i w_i (x_i - \mu_x) (y_i - \mu_y)
# = \sum_i w_i x_i y_i - (\sum_i w_i x_i) (\sum_j w_j y_j).
num1 = reducer(img1 * img2, kernel) * 2.0
den1 = reducer(tf.square(img1) + tf.square(img2), kernel)
c_s = (num1 - num0 + c_two) / (den1 - den0 + c_two)
# SSIM score is the product of the luminance and contrast-structure measures.
return luminance, c_s
filter_size = tf.constant(9, dtype='int32') # changed from 11 to 9 due
filter_sigma = tf.constant(1.5, dtype=img1.dtype)
shape1, shape2 = tf.shape_n([img1, img2])
checks = [
tf.Assert(tf.reduce_all(
tf.greater_equal(shape1[-3:-1], filter_size)),
[shape1, filter_size],
summarize=8),
tf.Assert(tf.reduce_all(
tf.greater_equal(shape2[-3:-1], filter_size)),
[shape2, filter_size],
summarize=8)
]
# Enforce the check to run before computation.
with tf.control_dependencies(checks):
img1 = tf.identity(img1)
# TODO(sjhwang): Try to cache kernels and compensation factor.
kernel = _fspecial_gauss(filter_size, filter_sigma)
kernel = tf.tile(kernel, multiples=[1, 1, shape1[-1], 1])
# The correct compensation factor is `1.0 - tf.reduce_sum(tf.square(kernel))`,
# but to match MATLAB implementation of MS-SSIM, we use 1.0 instead.
compensation = 1.
# TODO(sjhwang): Try FFT.
# TODO(sjhwang): Gaussian kernel is separable in space. Consider applying
# 1-by-n and n-by-1 Gaussain filters instead of an n-by-n filter.
luminance, c_s = _ssim_helper(img1, img2, max_val, kernel,
compensation)
# Average over the second and the third from the last: height, width.
axes = tf.constant([-3, -2], dtype='int32')
ssim_val = tf.reduce_mean(luminance * c_s, axes)
c_s = tf.reduce_mean(c_s, axes)
return ssim_val, c_s
def function_factory(self, loss, train_x, train_y):
"""A factory to create a function required by tfp.optimizer.lbfgs_minimize.
Args:
self [in]: an instance of `tf.keras.Model` or its subclasses.
loss [in]: a function with signature loss_value = loss(pred_y, true_y).
train_x [in]: the input part of training data.
train_y [in]: the output part of training data.
Returns:
A function that has a signature of:
loss_value, gradients = f(model_parameters).
"""
# obtain the shapes of all trainable parameters in the model
shapes = tf.shape_n(self.trainable_variables)
n_tensors = len(shapes)
# we'll use tf.dynamic_stitch and tf.dynamic_partition later, so we need to
# prepare required information first
count = 0
idx = [] # stitch indices
part = [] # partition indices
for i, shape in enumerate(shapes):
n = numpy.product(shape)
idx.append(
tf.reshape(tf.range(count, count + n, dtype=tf.int32), shape))
part.extend([i] * n)
count += n
part = tf.constant(part)
@tf.function
def assign_new_model_parameters(params_1d):
"""A function updating the model's parameters with a 1D tf.Tensor.
Args:
params_1d [in]: a 1D tf.Tensor representing the model's trainable parameters.
"""
params = tf.dynamic_partition(params_1d, part, n_tensors)
for i, (shape, param) in enumerate(zip(shapes, params)):
self.trainable_variables[i].assign(tf.reshape(param, shape))
# now create a function that will be returned by this factory
@tf.function
def f(params_1d):
"""A function that can be used by tfp.optimizer.lbfgs_minimize.
This function is created by function_factory.
Args:
params_1d [in]: a 1D tf.Tensor.
Returns:
A scalar loss and the gradients w.r.t. the `params_1d`.
"""
# use GradientTape so that we can calculate the gradient of loss w.r.t. parameters
with tf.GradientTape() as tape:
# update the parameters in the model
assign_new_model_parameters(params_1d)
# calculate the loss
loss_value = loss(self(train_x, training=True), train_y)
# calculate gradients and convert to 1D tf.Tensor
grads = tape.gradient(loss_value, self.trainable_variables)
grads = [
k if k is not None else tf.zeros_like(
self.trainable_variables[i]) for i, k in enumerate(grads)
]
grads = tf.dynamic_stitch(idx, grads)
# print out iteration & loss
f.iter.assign_add(1)
tf.print("Iter:", f.iter, "loss:", loss_value, "\t", end='\r')
return loss_value, grads
# store these information as members so we can use them outside the scope
f.iter = tf.Variable(0)
f.idx = idx
f.part = part
f.shapes = shapes
f.assign_new_model_parameters = assign_new_model_parameters
return f
a = tf.Variable([[1, 2, 4], [3, 4, 5]])
s = tf.shape(a)
sess = tf.Session()
sess.run(tf.global_variables_initializer())
out = sess.run(s)
print("shape:", out)
size = tf.size(a)
out = sess.run(size)
print("size:", out)
rank = tf.rank(a)
out = sess.run(rank)
print("rank:", out)
shape_n = tf.shape_n([a])
out = sess.run(shape_n)
print("shape_n:", out)
'''
5. 改变tensor的shape: tf.reshape
6. 扩展维度： tf.expand_dims
7. 去掉size为1的维度：tf.squeeze(input,axis=None,name=None,squeeze_dims=None)
参数：
* input: Tensor
* axis: a list of ints, 可选。要去掉的维度。
* name：
* squeeze_dims: 过时的参数，同axis。
8. Broadcasts parameters : tf.meshgrid()
'''
def graph_lbfgs(model, loss):
"""A factory to create a function required by tfp.optimizer.lbfgs_minimize.
Args:
model [in]: an instance of `tf.keras.Model` or its subclasses.
loss [in]: a function with signature loss_value = loss(pred_y, true_y).
Returns:
A function that has a signature of:
loss_value, gradients = f(model_parameters).
"""
# obtain the shapes of all trainable parameters in the model
shapes = tf.shape_n(model.trainable_variables)
n_tensors = len(shapes)
# we'll use tf.dynamic_stitch and tf.dynamic_partition later, so we need to
# prepare required information first
count = 0
idx = [] # stitch indices
part = [] # partition indices
start_time = time.time()
for i, shape in enumerate(shapes):
n = numpy.product(shape)
idx.append(
tf.reshape(tf.range(count, count + n, dtype=tf.int32), shape))
part.extend([i] * n)
count += n
part = tf.constant(part)
@tf.function
def assign_new_model_parameters(params_1d):
"""A function updating the model's parameters with a 1D tf.Tensor.
Args:
params_1d [in]: a 1D tf.Tensor representing the model's trainable parameters.
"""
params = tf.dynamic_partition(params_1d, part, n_tensors)
for i, (shape, param) in enumerate(zip(shapes, params)):
model.trainable_variables[i].assign(tf.reshape(param, shape))
# now create a function that will be returned by this factory
@tf.function
def f(params_1d):
"""A function that can be used by tfp.optimizer.lbfgs_minimize.
This function is created by function_factory.
Args:
params_1d [in]: a 1D tf.Tensor.
Returns:
A scalar loss and the gradients w.r.t. the `params_1d`.
"""
# use GradientTape so that we can calculate the gradient of loss w.r.t. parameters
with tf.GradientTape() as tape:
# update the parameters in the model
assign_new_model_parameters(params_1d)
# calculate the loss
loss_value = loss()
# calculate gradients and convert to 1D tf.Tensor
grads = tape.gradient(loss_value, model.trainable_variables)
grads = tf.dynamic_stitch(idx, grads)
# print out iteration & loss
f.iter.assign_add(1)
if f.iter % 300 == 0:
elapsed = tf.timestamp() - f.start_time
tf.print("Iter:", f.iter // 3, "loss:", loss_value, "time:",
elapsed)
f.start_time.assign(tf.timestamp())
# store loss value so we can retrieve later
tf.py_function(f.history.append, inp=[loss_value], Tout=[])
return loss_value, grads
# store these information as members so we can use them outside the scope
f.iter = tf.Variable(0)
f.idx = idx
f.part = part
f.shapes = shapes
f.assign_new_model_parameters = assign_new_model_parameters
f.history = []
f.start_time = tf.Variable(tf.timestamp())
return f
def _calculateRpnLoss(predRawScores, predBoxes, predRegressions, predAnchors,
mini_batch_size, gt_boxes, feature_h, feature_w):
"""Refactoring of code from calculateRpnLoss into another function for testing"""
num_classes = tf.shape_n(
gt_boxes)[0] - 1 # subtract one since background is not a class
iou_threshold_neg = s.DEF_IOU_THRESHOLD_TRAIN_NEG
iou_threshold_pos = s.DEF_IOU_THRESHOLD_TRAIN_POS
with easy_scope(name="proposal_layer_test"), tf.device("/cpu:0"):
labeled_boxes = iou_labeler(predBoxes, gt_boxes, iou_threshold_neg,
iou_threshold_pos)
# Sample boxes and raw scores for loss
posIdx, negIdx = tf.py_func(
lambda x: sampleBoxes(x, num_classes, mini_batch_size),
[labeled_boxes], [tf.int32, tf.int32],
stateful=False,
name="sampleBoxes")
# posIdx, negIdx = tf.py_func(sampleBoxes, [labeled_boxes, num_classes, mini_batch_size],
# tf.float32, stateful=False, name="sampleBoxes")
positive_raw_scores = tf.gather(predRawScores,
posIdx,
axis=0,
name="positive_raw_scores")
negative_raw_scores = tf.gather(predRawScores,
negIdx,
axis=0,
name="negative_raw_scores")
# There is no regression loss for negative examples. For the positives, we need
# to find the gt regression from anchor to gt boxes
positive_anchors = tf.gather(predAnchors,
posIdx,
axis=0,
name="positive_anchors")
positive_gt_boxes = tf.gather(gt_boxes,
tf.cast(tf.gather(
labeled_boxes[:, 4], posIdx),
dtype=tf.int32),
name="positive_gt_boxes")
positive_gt_regs = calculateRegressions(positive_anchors,
positive_gt_boxes,
axis=-1)
positive_raw_regressions = tf.gather(predRegressions,
posIdx,
axis=0,
name="positive_raw_regressions")
# Flatten regressions before passing into the huber loss function
flat_pred_regs = tf.reshape(positive_raw_regressions, [-1])
flat_gt_regs = tf.reshape(positive_gt_regs, [-1])
reg_loss = tf.losses.huber_loss(flat_pred_regs,
flat_gt_regs,
reduction=tf.losses.Reduction.NONE,
delta=1.0)
reg_loss = tf.reduce_sum(reg_loss)
# Class-agnostic log loss for positive examples
# Need to create a whole bunch of [0,1]s of the right length
num_pos = tf.shape(positive_raw_scores)[0]
cls_loss_pos = tf.nn.sparse_softmax_cross_entropy_with_logits(
labels=tf.ones([num_pos], dtype=tf.int32),
logits=positive_raw_scores)
cls_loss_pos = tf.reduce_sum(cls_loss_pos)
# Log-loss for the negative examples
num_neg = tf.shape(negative_raw_scores)[0]
cls_loss_neg = tf.nn.sparse_softmax_cross_entropy_with_logits(
labels=tf.zeros([num_neg], dtype=tf.int32),
logits=negative_raw_scores)
cls_loss_neg = tf.reduce_sum(cls_loss_neg)
# Adding up and normalizing the losses.
reg_loss /= (feature_h * feature_w) / 10.
cls_loss = (cls_loss_pos + cls_loss_neg) / mini_batch_size
total_loss = tf.add(reg_loss, cls_loss, name="total_loss")
return total_loss
