def __call__(self, tensors):
kl = concrete_binary_sample_kl(tensors["obj_pre_sigmoid"],
tensors["obj_log_odds"],
self.obj_concrete_temp,
self.prior_log_odds,
self.obj_concrete_temp)
batch_size = tf_shape(tensors["obj_pre_sigmoid"])[0]
return tf.reduce_sum(tf.reshape(kl, (batch_size, -1)), 1)
def body(step, stopping_sum, prev_state, running_recon, kl_loss,
running_digits, scale_ta, scale_kl_ta, scale_std_ta, shift_ta,
shift_kl_ta, shift_std_ta, attr_ta, attr_kl_ta, attr_std_ta,
z_pres_ta, z_pres_probs_ta, z_pres_kl_ta, vae_input_ta,
vae_output_ta, scale, shift, attr, z_pres):
if self.difference_air:
inp = (self._tensors["inp"] -
tf.reshape(running_recon,
(self.batch_size, *self.obs_shape)))
encoded_inp = self.image_encoder(inp, 0, self.is_training)
encoded_inp = tf.layers.flatten(encoded_inp)
else:
encoded_inp = self.encoded_inp
if self.complete_rnn_input:
rnn_input = tf.concat(
[encoded_inp, scale, shift, attr, z_pres], axis=1)
else:
rnn_input = encoded_inp
hidden_rep, next_state = self.cell(rnn_input, prev_state)
outputs = self.output_network(hidden_rep, 9, self.is_training)
(scale_mean, scale_log_std, shift_mean, shift_log_std,
z_pres_log_odds) = tf.split(outputs, [2, 2, 2, 2, 1], axis=1)
# --- scale ---
scale_std = tf.exp(scale_log_std)
scale_mean = self.apply_fixed_value("scale_mean", scale_mean)
scale_std = self.apply_fixed_value("scale_std", scale_std)
scale_logits, scale_kl = normal_vae(scale_mean, scale_std,
self.scale_prior_mean,
self.scale_prior_std)
scale_kl = tf.reduce_sum(scale_kl, axis=1, keepdims=True)
scale = tf.nn.sigmoid(tf.clip_by_value(scale_logits, -10, 10))
# --- shift ---
shift_std = tf.exp(shift_log_std)
shift_mean = self.apply_fixed_value("shift_mean", shift_mean)
shift_std = self.apply_fixed_value("shift_std", shift_std)
shift_logits, shift_kl = normal_vae(shift_mean, shift_std,
self.shift_prior_mean,
self.shift_prior_std)
shift_kl = tf.reduce_sum(shift_kl, axis=1, keepdims=True)
shift = tf.nn.tanh(tf.clip_by_value(shift_logits, -10, 10))
# --- Extract windows from scene ---
w, h = scale[:, 0:1], scale[:, 1:2]
x, y = shift[:, 0:1], shift[:, 1:2]
theta = tf.concat(
[w, tf.zeros_like(w), x,
tf.zeros_like(h), h, y], axis=1)
theta = tf.reshape(theta, (-1, 2, 3))
vae_input = transformer(self._tensors["inp"], theta,
self.object_shape)
# This is a necessary reshape, as the output of transformer will have unknown dims
vae_input = tf.reshape(
vae_input,
(self.batch_size, *self.object_shape, self.image_depth))
# --- Apply Object-level VAE (object encoder/object decoder) to windows ---
attr = self.object_encoder(vae_input, 2 * self.A, self.is_training)
attr_mean, attr_log_std = tf.split(attr, 2, axis=1)
attr_std = tf.exp(attr_log_std)
attr, attr_kl = normal_vae(attr_mean, attr_std,
self.attr_prior_mean,
self.attr_prior_std)
attr_kl = tf.reduce_sum(attr_kl, axis=1, keepdims=True)
vae_output = self.object_decoder(
attr,
self.object_shape[0] * self.object_shape[1] * self.image_depth,
self.is_training)
vae_output = tf.nn.sigmoid(tf.clip_by_value(vae_output, -10, 10))
# --- Place reconstructed objects in image ---
theta_inverse = tf.concat([
1. / w,
tf.zeros_like(w), -x / w,
tf.zeros_like(h), 1. / h, -y / h
],
axis=1)
theta_inverse = tf.reshape(theta_inverse, (-1, 2, 3))
vae_output_transformed = transformer(
tf.reshape(vae_output, (
self.batch_size,
*self.object_shape,
self.image_depth,
)), theta_inverse, self.obs_shape[:2])
vae_output_transformed = tf.reshape(vae_output_transformed, [
self.batch_size,
self.image_height * self.image_width * self.image_depth
])
# --- z_pres ---
if self.run_all_time_steps:
z_pres = tf.ones_like(z_pres_log_odds)
z_pres_prob = tf.ones_like(z_pres_log_odds)
z_pres_kl = tf.zeros_like(z_pres_log_odds)
else:
z_pres_log_odds = tf.clip_by_value(z_pres_log_odds, -10, 10)
z_pres_pre_sigmoid = concrete_binary_pre_sigmoid_sample(
z_pres_log_odds, self.z_pres_temperature)
z_pres = tf.nn.sigmoid(z_pres_pre_sigmoid)
z_pres = (self.float_is_training * z_pres +
(1 - self.float_is_training) * tf.round(z_pres))
z_pres_prob = tf.nn.sigmoid(z_pres_log_odds)
z_pres_kl = concrete_binary_sample_kl(
z_pres_pre_sigmoid,
z_pres_log_odds,
self.z_pres_temperature,
self.z_pres_prior_log_odds,
self.z_pres_temperature,
)
stopping_sum += (1.0 - z_pres)
alive = tf.less(stopping_sum, self.stopping_threshold)
running_digits += tf.to_int32(alive)
# --- adjust reconstruction ---
running_recon += tf.where(
tf.tile(alive, (1, vae_output_transformed.shape[1])),
z_pres * vae_output_transformed, tf.zeros_like(running_recon))
# --- add kl to loss ---
kl_loss += tf.where(alive, scale_kl, tf.zeros_like(kl_loss))
kl_loss += tf.where(alive, shift_kl, tf.zeros_like(kl_loss))
kl_loss += tf.where(alive, attr_kl, tf.zeros_like(kl_loss))
kl_loss += tf.where(alive, z_pres_kl, tf.zeros_like(kl_loss))
# --- record values ---
scale_ta = scale_ta.write(scale_ta.size(), scale)
scale_kl_ta = scale_kl_ta.write(scale_kl_ta.size(), scale_kl)
scale_std_ta = scale_std_ta.write(scale_std_ta.size(), scale_std)
shift_ta = shift_ta.write(shift_ta.size(), shift)
shift_kl_ta = shift_kl_ta.write(shift_kl_ta.size(), shift_kl)
shift_std_ta = shift_std_ta.write(shift_std_ta.size(), shift_std)
attr_ta = attr_ta.write(attr_ta.size(), attr)
attr_kl_ta = attr_kl_ta.write(attr_kl_ta.size(), attr_kl)
attr_std_ta = attr_std_ta.write(attr_std_ta.size(), attr_std)
vae_input_ta = vae_input_ta.write(vae_input_ta.size(),
tf.layers.flatten(vae_input))
vae_output_ta = vae_output_ta.write(vae_output_ta.size(),
vae_output)
z_pres_ta = z_pres_ta.write(z_pres_ta.size(), z_pres)
z_pres_probs_ta = z_pres_probs_ta.write(z_pres_probs_ta.size(),
z_pres_prob)
z_pres_kl_ta = z_pres_kl_ta.write(z_pres_kl_ta.size(), z_pres_kl)
return (
step + 1,
stopping_sum,
next_state,
running_recon,
kl_loss,
running_digits,
scale_ta,
scale_kl_ta,
scale_std_ta,
shift_ta,
shift_kl_ta,
shift_std_ta,
attr_ta,
attr_kl_ta,
attr_std_ta,
z_pres_ta,
z_pres_probs_ta,
z_pres_kl_ta,
vae_input_ta,
vae_output_ta,
scale,
shift,
attr,
z_pres,
)
def _compute_obj_kl(self, tensors, existing_objects=None):
# --- compute obj_kl ---
obj_pre_sigmoid = tensors["obj_pre_sigmoid"]
obj_log_odds = tensors["obj_log_odds"]
obj_prob = tensors["obj_prob"]
obj = tensors["obj"]
batch_size, n_objects, _ = tf_shape(obj)
max_n_objects = n_objects
if existing_objects is not None:
_, n_existing_objects, _ = tf_shape(existing_objects)
existing_objects = tf.reshape(existing_objects, (batch_size, n_existing_objects))
max_n_objects += n_existing_objects
count_support = tf.range(max_n_objects+1, dtype=tf.float32)
if self.count_prior_dist is not None:
if self.count_prior_dist is not None:
assert len(self.count_prior_dist) == (max_n_objects + 1)
count_distribution = tf.constant(self.count_prior_dist, dtype=tf.float32)
else:
count_prior_prob = tf.nn.sigmoid(self.count_prior_log_odds)
count_distribution = (1 - count_prior_prob) * (count_prior_prob ** count_support)
normalizer = tf.reduce_sum(count_distribution)
count_distribution = count_distribution / tf.maximum(normalizer, 1e-6)
count_distribution = tf.tile(count_distribution[None, :], (batch_size, 1))
if existing_objects is not None:
count_so_far = tf.reduce_sum(tf.round(existing_objects), axis=1, keepdims=True)
count_distribution = (
count_distribution
* tf_binomial_coefficient(count_support, count_so_far)
* tf_binomial_coefficient(max_n_objects - count_support, n_existing_objects - count_so_far)
)
normalizer = tf.reduce_sum(count_distribution, axis=1, keepdims=True)
count_distribution = count_distribution / tf.maximum(normalizer, 1e-6)
else:
count_so_far = tf.zeros((batch_size, 1), dtype=tf.float32)
obj_kl = []
for i in range(n_objects):
p_z_given_Cz_raw = (count_support[None, :] - count_so_far) / (max_n_objects - i)
p_z_given_Cz = tf.clip_by_value(p_z_given_Cz_raw, 0.0, 1.0)
# Doing this instead of 1 - p_z_given_Cz seems to be more numerically stable.
inv_p_z_given_Cz_raw = (max_n_objects - i - count_support[None, :] + count_so_far) / (max_n_objects - i)
inv_p_z_given_Cz = tf.clip_by_value(inv_p_z_given_Cz_raw, 0.0, 1.0)
p_z = tf.reduce_sum(count_distribution * p_z_given_Cz, axis=1, keepdims=True)
if self.use_concrete_kl:
prior_log_odds = tf_safe_log(p_z) - tf_safe_log(1-p_z)
_obj_kl = concrete_binary_sample_kl(
obj_pre_sigmoid[:, i, :],
obj_log_odds[:, i, :], self.obj_concrete_temp,
prior_log_odds, self.obj_concrete_temp,
)
else:
prob = obj_prob[:, i, :]
_obj_kl = (
prob * (tf_safe_log(prob) - tf_safe_log(p_z))
+ (1-prob) * (tf_safe_log(1-prob) - tf_safe_log(1-p_z))
)
obj_kl.append(_obj_kl)
sample = tf.to_float(obj[:, i, :] > 0.5)
mult = sample * p_z_given_Cz + (1-sample) * inv_p_z_given_Cz
raw_count_distribution = mult * count_distribution
normalizer = tf.reduce_sum(raw_count_distribution, axis=1, keepdims=True)
normalizer = tf.maximum(normalizer, 1e-6)
# invalid = tf.logical_and(p_z_given_Cz_raw > 1, count_distribution > 1e-8)
# float_invalid = tf.cast(invalid, tf.float32)
# diagnostic = tf.stack(
# [float_invalid, p_z_given_Cz, count_distribution, mult, raw_count_distribution], axis=-1)
# assert_op = tf.Assert(
# tf.reduce_all(tf.logical_not(invalid)),
# [invalid, diagnostic, count_so_far, sample, tf.constant(i, dtype=tf.float32)],
# summarize=100000)
count_distribution = raw_count_distribution / normalizer
count_so_far += sample
# this avoids buildup of inaccuracies that can cause problems in computing p_z_given_Cz_raw
count_so_far = tf.round(count_so_far)
obj_kl = tf.reshape(tf.concat(obj_kl, axis=1), (batch_size, n_objects, 1))
return obj_kl
def _compute_obj_kl(self, tensors):
# --- compute obj_kl ---
count_support = tf.range(self.HWB + 1, dtype=tf.float32)
if self.count_prior_dist is not None:
if self.count_prior_dist is not None:
assert len(self.count_prior_dist) == (self.HWB + 1)
count_distribution = tf.constant(self.count_prior_dist,
dtype=tf.float32)
else:
count_prior_prob = tf.nn.sigmoid(self.count_prior_log_odds)
count_distribution = (1 - count_prior_prob) * (count_prior_prob**
count_support)
normalizer = tf.reduce_sum(count_distribution)
count_distribution = count_distribution / normalizer
count_distribution = tf.tile(count_distribution[None, :],
(self.batch_size, 1))
count_so_far = tf.zeros((self.batch_size, 1), dtype=tf.float32)
i = 0
obj_kl = []
for h, w, b in itertools.product(range(self.H), range(self.W),
range(self.B)):
p_z_given_Cz = tf.maximum(count_support[None, :] - count_so_far,
0) / (self.HWB - i)
# Reshape for batch matmul
_count_distribution = count_distribution[:, None, :]
_p_z_given_Cz = p_z_given_Cz[:, :, None]
p_z = tf.matmul(_count_distribution, _p_z_given_Cz)[:, :, 0]
if self.use_concrete_kl:
prior_log_odds = tf_safe_log(p_z) - tf_safe_log(1 - p_z)
_obj_kl = concrete_binary_sample_kl(
tensors["obj_pre_sigmoid"][:, h, w, b, :],
tensors["obj_log_odds"][:, h, w, b, :],
self.obj_concrete_temp,
prior_log_odds,
self.obj_concrete_temp,
)
else:
prob = tensors["obj_prob"][:, h, w, b, :]
_obj_kl = (prob * (tf_safe_log(prob) - tf_safe_log(p_z)) +
(1 - prob) *
(tf_safe_log(1 - prob) - tf_safe_log(1 - p_z)))
obj_kl.append(_obj_kl)
sample = tf.to_float(tensors["obj"][:, h, w, b, :] > 0.5)
mult = sample * p_z_given_Cz + (1 - sample) * (1 - p_z_given_Cz)
count_distribution = mult * count_distribution
normalizer = tf.reduce_sum(count_distribution,
axis=1,
keepdims=True)
normalizer = tf.maximum(normalizer, 1e-6)
count_distribution = count_distribution / normalizer
count_so_far += sample
i += 1
obj_kl = tf.reshape(tf.concat(obj_kl, axis=1),
(self.batch_size, self.H, self.W, self.B, 1))
return obj_kl
def compute_kl(self, tensors, prior=None):
simple_obj_kl = prior is not None
if prior is None:
prior = self._independent_prior()
# --- box ---
cell_y_kl = tensors["cell_y_logit_dist"].kl_divergence(
prior["cell_y_logit_dist"])
cell_x_kl = tensors["cell_x_logit_dist"].kl_divergence(
prior["cell_x_logit_dist"])
height_kl = tensors["height_logit_dist"].kl_divergence(
prior["height_logit_dist"])
width_kl = tensors["width_logit_dist"].kl_divergence(
prior["width_logit_dist"])
if "cell_y" in self.no_gradient:
cell_y_kl = tf.stop_gradient(cell_y_kl)
if "cell_x" in self.no_gradient:
cell_x_kl = tf.stop_gradient(cell_x_kl)
if "height" in self.no_gradient:
height_kl = tf.stop_gradient(height_kl)
if "width" in self.no_gradient:
width_kl = tf.stop_gradient(width_kl)
box_kl = tf.concat([cell_y_kl, cell_x_kl, height_kl, width_kl],
axis=-1)
# --- attr ---
attr_kl = tensors["attr_dist"].kl_divergence(prior["attr_dist"])
if "attr" in self.no_gradient:
attr_kl = tf.stop_gradient(attr_kl)
# --- z ---
z_kl = tensors["z_logit_dist"].kl_divergence(prior["z_logit_dist"])
if "z" in self.no_gradient:
z_kl = tf.stop_gradient(z_kl)
if "z" in self.fixed_values:
z_kl = tf.zeros_like(z_kl)
# --- obj ---
if simple_obj_kl:
obj_kl = concrete_binary_sample_kl(
tensors["obj_pre_sigmoid"],
tensors["obj_log_odds"],
self.obj_concrete_temp,
prior["obj_log_odds"],
self.obj_concrete_temp,
)
else:
obj_kl = self._compute_obj_kl(tensors)
if "obj" in self.no_gradient:
obj_kl = tf.stop_gradient(obj_kl)
return dict(
cell_y_kl=cell_y_kl,
cell_x_kl=cell_x_kl,
height_kl=height_kl,
width_kl=width_kl,
box_kl=box_kl,
z_kl=z_kl,
attr_kl=attr_kl,
obj_kl=obj_kl,
)
def _compute_obj_kl(self, tensors, existing_objects=None):
# --- compute obj_kl ---
obj_pre_sigmoid = tensors["obj_pre_sigmoid"]
obj_log_odds = tensors["obj_log_odds"]
obj_prob = tensors["obj_prob"]
obj = tensors["obj"]
batch_size, n_objects, _ = tf_shape(obj)
max_n_objects = n_objects
if existing_objects is not None:
_, n_existing_objects, _ = tf_shape(existing_objects)
existing_objects = tf.reshape(existing_objects,
(batch_size, n_existing_objects))
max_n_objects += n_existing_objects
count_support = tf.range(max_n_objects + 1, dtype=tf.float32)
if self.count_prior_dist is not None:
if self.count_prior_dist is not None:
assert len(self.count_prior_dist) == (max_n_objects + 1)
count_distribution = tf.constant(self.count_prior_dist,
dtype=tf.float32)
else:
count_prior_prob = tf.nn.sigmoid(self.count_prior_log_odds)
count_distribution = (1 - count_prior_prob) * (count_prior_prob**
count_support)
normalizer = tf.reduce_sum(count_distribution)
count_distribution = count_distribution / tf.maximum(normalizer, 1e-6)
count_distribution = tf.tile(count_distribution[None, :],
(batch_size, 1))
if existing_objects is not None:
count_so_far = tf.reduce_sum(tf.round(existing_objects),
axis=1,
keepdims=True)
count_distribution = (
count_distribution *
tf_binomial_coefficient(count_support, count_so_far) *
tf_binomial_coefficient(max_n_objects - count_support,
n_existing_objects - count_so_far))
normalizer = tf.reduce_sum(count_distribution,
axis=1,
keepdims=True)
count_distribution = count_distribution / tf.maximum(
normalizer, 1e-6)
else:
count_so_far = tf.zeros((batch_size, 1), dtype=tf.float32)
obj_kl = []
for i in range(n_objects):
p_z_given_Cz = tf.maximum(count_support[None, :] - count_so_far,
0) / (max_n_objects - i)
# Reshape for batch matmul
_count_distribution = count_distribution[:, None, :]
_p_z_given_Cz = p_z_given_Cz[:, :, None]
p_z = tf.matmul(_count_distribution, _p_z_given_Cz)[:, :, 0]
if self.use_concrete_kl:
prior_log_odds = tf_safe_log(p_z) - tf_safe_log(1 - p_z)
_obj_kl = concrete_binary_sample_kl(
obj_pre_sigmoid[:, i, :],
obj_log_odds[:, i, :],
self.obj_concrete_temp,
prior_log_odds,
self.obj_concrete_temp,
)
else:
prob = obj_prob[:, i, :]
_obj_kl = (prob * (tf_safe_log(prob) - tf_safe_log(p_z)) +
(1 - prob) *
(tf_safe_log(1 - prob) - tf_safe_log(1 - p_z)))
obj_kl.append(_obj_kl)
sample = tf.to_float(obj[:, i, :] > 0.5)
mult = sample * p_z_given_Cz + (1 - sample) * (1 - p_z_given_Cz)
count_distribution = mult * count_distribution
normalizer = tf.reduce_sum(count_distribution,
axis=1,
keepdims=True)
normalizer = tf.maximum(normalizer, 1e-6)
count_distribution = count_distribution / normalizer
count_so_far += sample
obj_kl = tf.reshape(tf.concat(obj_kl, axis=1),
(batch_size, n_objects, 1))
return obj_kl