def _get_coords_for_joint(joint_idx, parent_idx, child_angle_idx,
coords):
if parent_idx is None: # joint_idx should be 0
coords[joint_idx] = K.zeros(base_shape[:-2] + [3, 1])
parent_bone = K.constant(
np.concatenate([
np.ones(base_shape),
np.zeros(base_shape),
np.zeros(base_shape)
],
axis=-2))
else:
parent_bone = coords[parent_idx] - coords[joint_idx]
parent_bone_norm = K.sqrt(
K.sum(K.square(parent_bone), axis=-2, keepdims=True) +
K.epsilon())
parent_bone = parent_bone / parent_bone_norm
for child_idx in body_graph[joint_idx]:
child_bone = tf.matmul(rotmat_list[child_angle_idx],
parent_bone)
child_bone_idx = bone_idcs[(joint_idx, child_idx)]
child_bone = child_bone * K.reshape(
bone_len_list[child_bone_idx],
(child_bone.shape[0], 1, 1, 1))
coords[child_idx] = child_bone + coords[joint_idx]
child_angle_idx += 1
for child_idx in body_graph[joint_idx]:
child_angle_idx, coords = _get_coords_for_joint(
child_idx, joint_idx, child_angle_idx, coords)
return child_angle_idx, coords
def _get_bone_len(arg):
bone_list = tf.unstack(arg[:, :, 0, :], axis=1)
bones = [
bone_list[j] - bone_list[i]
for i, j in zip(members_from, members_to)
]
bones = K.stack(bones, axis=1)
return K.sqrt(K.sum(K.square(bones), axis=-1) + K.epsilon())
def __call__(self, x):
regularization = 0.
if self.l1:
# X=K.eval(x)
diff = x[1:] - x[:-1]
regularization += K.mean(K.sqrt(diff**2 + 0.000001))
return regularization * self.l1
def _get_avg_bone_len(arg):
bone_list = tf.unstack(arg[:, :, 0, :], axis=1)
bones = [
bone_list[j] - bone_list[i]
for i, j in zip(members_from, members_to)
]
bones = K.expand_dims(K.stack(bones, axis=1), axis=2)
bone_len = K.sqrt(
K.sum(K.square(bones), axis=-1, keepdims=True) + K.epsilon())
return K.mean(bone_len, axis=1, keepdims=True)
def call(self, inputs, **kwargs):
input_shape = K.int_shape(inputs)
tensor_input_shape = K.shape(inputs)
# Prepare broadcasting shape.
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] // self.groups
broadcast_shape.insert(1, self.groups)
reshape_group_shape = K.shape(inputs)
group_axes = [reshape_group_shape[i] for i in range(len(input_shape))]
group_axes[self.axis] = input_shape[self.axis] // self.groups
group_axes.insert(1, self.groups)
# reshape inputs to new group shape
group_shape = [group_axes[0], self.groups] + group_axes[2:]
group_shape = K.stack(group_shape)
inputs = K.reshape(inputs, group_shape)
group_reduction_axes = list(range(len(group_axes)))
mean, variance = _moments(inputs,
group_reduction_axes[2:],
keep_dims=True)
inputs = (inputs - mean) / (K.sqrt(variance + self.epsilon))
# prepare broadcast shape
inputs = K.reshape(inputs, group_shape)
outputs = inputs
# In this case we must explicitly broadcast all parameters.
if self.scale:
broadcast_gamma = K.reshape(self.gamma, broadcast_shape)
outputs = outputs * broadcast_gamma
if self.center:
broadcast_beta = K.reshape(self.beta, broadcast_shape)
outputs = outputs + broadcast_beta
# finally we reshape the output back to the input shape
outputs = K.reshape(outputs, tensor_input_shape)
return outputs
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 edm(x, y=None):
with K.name_scope('edm'):
y = x if y is None else y
x = K.expand_dims(x, axis=1)
y = K.expand_dims(y, axis=2)
return K.sqrt(K.sum(K.square(x - y), axis=-1) + K.epsilon())
def __call__(self, p):
p *= K.cast(p >= 0., K.floatx())
return p / (K.epsilon() +
K.sqrt(K.sum(K.square(p), axis=self.axis, keepdims=True)))
def rmse(y_true, y_pred):
return backend.sqrt(backend.mean(backend.square(y_pred - y_true), axis=-1))