def style_loss(style, gen):
assert K.ndim(style) == 3
assert K.ndim(gen) == 3
S = gram_matrix(style)
G = gram_matrix(gen)
channels = 3
size = img_h * img_w
# Euclidean distance of the gram matrices multiplied by the constant
return K.sum(K.square(S - G)) / (4. * (channels**2) * (size**2))
def total_variation_loss(x):
assert K.ndim(x) == 4
if K.image_data_format() == 'channels_first':
a = K.square(x[:, :, :img_h - 1, :img_w - 1] - x[:, :, 1:, :img_w - 1])
b = K.square(x[:, :, :img_h - 1, :img_w - 1] - x[:, :, :img_h - 1, 1:])
else:
# Move the image pixel by pixel, and calculate the variance
a = K.square(x[:, :img_h - 1, :img_w - 1, :] - x[:, 1:, :img_w - 1, :])
b = K.square(x[:, :img_h - 1, :img_w - 1, :] - x[:, :img_h - 1, 1:, :])
return K.sum(K.pow(a + b, 1.25))
def gram_matrix(x):
assert K.ndim(x) == 3
if K.image_data_format() == 'channels_first':
features = K.flatten(x)
else:
features = K.batch_flatten(K.permute_dimensions(x, (2, 0, 1)))
# Dot product of the flattened feature map and the transpose of the
# flattened feature map
gram = K.dot(features, K.transpose(features))
return gram
def call(self, x, mask=None):
if self.mode == 0 or self.mode == 2:
assert self.built, 'Layer must be built before being called'
input_shape = K.int_shape(x)
reduction_axes = list(range(len(input_shape)))
del reduction_axes[self.axis]
broadcast_shape = [1] * len(input_shape)
broadcast_shape[self.axis] = input_shape[self.axis]
mean_batch, var_batch = _moments(x,
reduction_axes,
shift=None,
keep_dims=False)
std_batch = (K.sqrt(var_batch + self.epsilon))
r_max_value = K.get_value(self.r_max)
r = std_batch / (K.sqrt(self.running_std + self.epsilon))
r = K.stop_gradient(K.clip(r, 1 / r_max_value, r_max_value))
d_max_value = K.get_value(self.d_max)
d = (mean_batch - self.running_mean) / K.sqrt(self.running_std +
self.epsilon)
d = K.stop_gradient(K.clip(d, -d_max_value, d_max_value))
if sorted(reduction_axes) == range(K.ndim(x))[:-1]:
x_normed_batch = (x - mean_batch) / std_batch
x_normed = (x_normed_batch * r + d) * self.gamma + self.beta
else:
# need broadcasting
broadcast_mean = K.reshape(mean_batch, broadcast_shape)
broadcast_std = K.reshape(std_batch, broadcast_shape)
broadcast_r = K.reshape(r, broadcast_shape)
broadcast_d = K.reshape(d, broadcast_shape)
broadcast_beta = K.reshape(self.beta, broadcast_shape)
broadcast_gamma = K.reshape(self.gamma, broadcast_shape)
x_normed_batch = (x - broadcast_mean) / broadcast_std
x_normed = (x_normed_batch * broadcast_r +
broadcast_d) * broadcast_gamma + broadcast_beta
# explicit update to moving mean and standard deviation
self.add_update([
K.moving_average_update(self.running_mean, mean_batch,
self.momentum),
K.moving_average_update(self.running_std, std_batch**2,
self.momentum)
], x)
# update r_max and d_max
r_val = self.r_max_value / (
1 + (self.r_max_value - 1) * K.exp(-self.t))
d_val = self.d_max_value / (1 + (
(self.d_max_value / 1e-3) - 1) * K.exp(-(2 * self.t)))
self.add_update([
K.update(self.r_max, r_val),
K.update(self.d_max, d_val),
K.update_add(self.t, K.variable(np.array([self.t_delta])))
], x)
if self.mode == 0:
if sorted(reduction_axes) == range(K.ndim(x))[:-1]:
x_normed_running = K.batch_normalization(
x,
self.running_mean,
self.running_std,
self.beta,
self.gamma,
epsilon=self.epsilon)
else:
# need broadcasting
broadcast_running_mean = K.reshape(self.running_mean,
broadcast_shape)
broadcast_running_std = K.reshape(self.running_std,
broadcast_shape)
broadcast_beta = K.reshape(self.beta, broadcast_shape)
broadcast_gamma = K.reshape(self.gamma, broadcast_shape)
x_normed_running = K.batch_normalization(
x,
broadcast_running_mean,
broadcast_running_std,
broadcast_beta,
broadcast_gamma,
epsilon=self.epsilon)
# pick the normalized form of x corresponding to the training phase
# for batch renormalization, inference time remains same as batchnorm
x_normed = K.in_train_phase(x_normed, x_normed_running)
elif self.mode == 1:
# sample-wise normalization
m = K.mean(x, axis=self.axis, keepdims=True)
std = K.sqrt(
K.var(x, axis=self.axis, keepdims=True) + self.epsilon)
x_normed_batch = (x - m) / (std + self.epsilon)
r_max_value = K.get_value(self.r_max)
r = std / (self.running_std + self.epsilon)
r = K.stop_gradient(K.clip(r, 1 / r_max_value, r_max_value))
d_max_value = K.get_value(self.d_max)
d = (m - self.running_mean) / (self.running_std + self.epsilon)
d = K.stop_gradient(K.clip(d, -d_max_value, d_max_value))
x_normed = ((x_normed_batch * r) + d) * self.gamma + self.beta
# update r_max and d_max
t_val = K.get_value(self.t)
r_val = self.r_max_value / (
1 + (self.r_max_value - 1) * np.exp(-t_val))
d_val = self.d_max_value / (1 + (
(self.d_max_value / 1e-3) - 1) * np.exp(-(2 * t_val)))
t_val += float(self.t_delta)
self.add_update([
K.update(self.r_max, r_val),
K.update(self.d_max, d_val),
K.update(self.t, t_val)
], x)
return x_normed
def content_loss(content, gen):
assert K.ndim(content) == 3
assert K.ndim(gen) == 3
# Euclidean distance
return K.sum(K.square(gen - content))