def step_fn(batch):
with tf.GradientTape(persistent=True) as tape:
##### Forward pass
known = stack_known(batch)
mean, logvar = self.encode(known)
z = self.reparameterize(mean, logvar)
sample_logit = self.decode(z)
sample = tf.nn.sigmoid(sample_logit)
output = {
'predicted_occ': sample,
'predicted_free': 1 - sample
}
metrics = nn.calc_metrics(output, batch)
#### vae loss
vae_loss = compute_vae_loss(z,
mean,
logvar,
sample_logit,
labels=batch['gt_occ'])
metrics['loss/vae'] = vae_loss
### gan loss
fake_occ = tf.cast(sample_logit > 0, tf.float32)
real_pair_est = self.discriminate(known, batch['gt_occ'])
fake_pair_est = self.discriminate(known, fake_occ)
gan_loss_g = 10000 * (1 + tf.reduce_mean(-fake_pair_est))
gan_loss_d_no_gp = 1 + tf.reduce_mean(fake_pair_est -
real_pair_est)
# gradient penalty
gp = self.gradient_penalty(known, batch['gt_occ'], fake_occ)
gan_loss_d = gan_loss_d_no_gp + gp
metrics['loss/gan_g'] = gan_loss_g
metrics['loss/gan_d'] = gan_loss_d
metrics['loss/gan_gp'] = gp
metrics['loss/gan_d_no_gp'] = gan_loss_d_no_gp
### apply
generator_loss = vae_loss + gan_loss_g
dis_loss = gan_loss_d
vae_variables = self.encoder.trainable_variables + self.generator.trainable_variables
vae_gradients = tape.gradient(generator_loss, vae_variables)
clipped_vae_gradients = [
tf.clip_by_value(g, -1e6, 1e6) for g in vae_gradients
]
self.opt.apply_gradients(
list(zip(clipped_vae_gradients, vae_variables)))
dis_variables = self.discriminator.trainable_variables
dis_gradients = tape.gradient(dis_loss, dis_variables)
clipped_dis_gradients = [
tf.clip_by_value(g, -1e6, 1e6) for g in dis_gradients
]
self.gan_opt.apply_gradients(
list(zip(clipped_dis_gradients, dis_variables)))
return generator_loss, metrics
def step_fn(batch):
with tf.GradientTape() as tape:
ae_features = self.encoder(self.prep_ae_input(batch))
cvcnn = self.cvcnn(
self.prep_cvcnn_inputs(batch['conditioned_occ'],
ae_features))
ae_output = self.decoder(ae_features)
p_occ = tf.nn.sigmoid(cvcnn['p_occ_logits'])
output = {'predicted_occ': p_occ, 'predicted_free': 1 - p_occ}
metrics = nn.calc_metrics(output, batch)
cross_ent = tf.nn.sigmoid_cross_entropy_with_logits(
logits=cvcnn['p_occ_logits'], labels=batch['gt_occ'])
vcnn_loss = nn.reduce_sum_batch(cross_ent)
ae_loss = nn.reduce_sum_batch(
tf.nn.sigmoid_cross_entropy_with_logits(
logits=ae_output, labels=batch['gt_occ']))
metrics['loss/aux_loss'] = ae_loss
loss = vcnn_loss + 0.1 * ae_loss
variables = self.trainable_variables
gradients = tape.gradient(loss, variables)
clipped_gradients = [
tf.clip_by_value(g, -1e6, 1e6) for g in gradients
]
self.opt.apply_gradients(
list(zip(clipped_gradients, variables)))
return loss, metrics
def step_fn(batch):
with tf.GradientTape() as tape:
output = self(batch)
metrics = nn.calc_metrics(output, batch)
if self.params['loss'] == 'cross_entropy':
cross_ent = tf.keras.losses.binary_crossentropy(
batch['gt_occ'], output['predicted_occ'])
loss = nn.reduce_sum_batch(cross_ent)
if self.params['stacknet_version'] == 'v4':
ae_loss = tf.reduce_sum(
tf.keras.losses.binary_crossentropy(
batch['gt_occ'], output['aux_occ']))
metrics['loss/aux_loss'] = ae_loss
metrics['loss/vcnn_loss'] = loss
loss = loss + 0.1 * ae_loss
variables = self.model.trainable_variables
gradients = tape.gradient(loss, variables)
clipped_gradients = [
tf.clip_by_value(g, -1e6, 1e6) for g in gradients
]
self.opt.apply_gradients(
list(zip(clipped_gradients, variables)))
return loss, metrics
def step_fn(batch):
with tf.GradientTape() as tape:
output_logits = self(batch, training=True)
x = tf.nn.sigmoid(output_logits['predicted_occ'])
output = {'predicted_occ': x, 'predicted_free': 1 - x}
metrics = nn.calc_metrics(output, batch)
if self.params['loss'] == 'mse':
loss = self.mse_loss(metrics)
elif self.params['loss'] == 'cross_entropy':
cross_ent = tf.nn.sigmoid_cross_entropy_with_logits(
logits=output_logits['predicted_occ'],
labels=batch['gt_occ'])
loss = nn.reduce_sum_batch(cross_ent)
variables = self.trainable_variables
gradients = tape.gradient(loss, variables)
self.opt.apply_gradients(list(zip(gradients, variables)))
metrics.update(self.get_insights(variables, gradients))
return loss, metrics
def step_fn_multiloss(batch):
with tf.GradientTape() as tape:
output_logits = self.model(self.prep_input(batch))
sample = tf.nn.sigmoid(output_logits)
output = {
'predicted_occ': sample,
'predicted_free': 1 - sample
}
metrics = nn.calc_metrics(output, batch)
cross_ent = tf.nn.sigmoid_cross_entropy_with_logits(
logits=output_logits, labels=batch['gt_occ'])
loss = nn.reduce_sum_batch(cross_ent)
metrics['loss/0_step'] = loss
m = metrics
for i in range(5):
# Multistep part
# if True:
# i = 0
b = {
'conditioned_occ': tf.cast(output_logits > 0,
tf.float32)
}
output_logits = self.model(b)
cross_ent = tf.nn.sigmoid_cross_entropy_with_logits(
logits=output_logits, labels=batch['gt_occ'])
step_loss = nn.reduce_sum_batch(cross_ent)
metrics['loss/{}_step'.format(i + 1)] = step_loss
# loss = loss + step_loss
# loss = tf.cond(metrics['pred|gt/p(predicted_occ|gt_occ)'] > 0.95,
# lambda: tf.add(step_loss, loss),
# lambda: loss)
variables = self.model.trainable_variables
gradients = tape.gradient(loss, variables)
clipped_gradients = [
tf.clip_by_value(g, -1e6, 1e6) for g in gradients
]
self.opt.apply_gradients(
list(zip(clipped_gradients, variables)))
return loss, metrics
def step_fn(batch):
with tf.GradientTape() as tape:
output = self(batch, training=True)
metrics = nn.calc_metrics(output, batch)
if self.params['loss'] == 'cross_entropy':
loss = tf.reduce_sum(
tf.keras.losses.binary_crossentropy(
batch['gt_occ'], output['predicted_occ']))
elif self.params['loss'] == 'mse':
loss = self.mse_loss(metrics)
variables = self.trainable_variables
gradients = tape.gradient(loss, variables)
self.opt.apply_gradients(list(zip(gradients, variables)))
return loss, metrics
def step_fn_multiloss(batch):
with tf.GradientTape() as tape:
ae_features = self.encoder(self.prep_ae_input(batch))
ae_output = self.decoder(ae_features)
ae_loss = nn.reduce_sum_batch(
tf.nn.sigmoid_cross_entropy_with_logits(
logits=ae_output, labels=batch['gt_occ']))
cvcnn_inp = self.prep_cvcnn_inputs(batch['conditioned_occ'],
ae_features)
loss = ae_loss
for i in range(6):
cvcnn = self.cvcnn(cvcnn_inp)
if i == 0:
p_occ = tf.nn.sigmoid(cvcnn['p_occ_logits'])
output = {
'predicted_occ': p_occ,
'predicted_free': 1 - p_occ
}
metrics = nn.calc_metrics(output, batch)
cross_ent = tf.nn.sigmoid_cross_entropy_with_logits(
logits=cvcnn['p_occ_logits'], labels=batch['gt_occ'])
step_loss = nn.reduce_sum_batch(cross_ent)
loss = loss + step_loss
metrics['loss/{}_step'.format(i)] = step_loss
cvcnn_inp = self.prep_cvcnn_inputs(
tf.cast(cvcnn['p_occ_logits'] > 0, tf.float32),
ae_features)
metrics['loss/aux_loss'] = ae_loss
variables = self.trainable_variables
gradients = tape.gradient(loss, variables)
clipped_gradients = [
tf.clip_by_value(g, -1e6, 1e6) for g in gradients
]
self.opt.apply_gradients(
list(zip(clipped_gradients, variables)))
return loss, metrics
def train_step(self, batch):
with tf.GradientTape() as tape:
known = stack_known(batch)
mean, logvar = self.encode(known)
true_angle = batch['angle']
z = self.reparameterize(mean, logvar)
z_corrected = self.replace_true_angle(z, true_angle)
sample_logit = self.decode(z_corrected)
z_f, sampled_angle = self.split_angle(z)
mean_f, mean_angle = self.split_angle(mean)
logvar_f, logvar_angle = self.split_angle(logvar)
vae_loss = compute_vae_loss(z_f,
mean_f,
logvar_f,
sample_logit,
labels=batch['gt_occ'])
angle_loss = compute_angle_loss(true_angle, mean_angle,
logvar_angle)
# angle_loss = compute_abs_angle_loss(true_angle, mean_angle)
loss = vae_loss + angle_loss
vae_variables = self.encoder.trainable_variables + self.generator.trainable_variables
gradients = tape.gradient(loss, vae_variables)
self.optimizer.apply_gradients(list(zip(gradients, vae_variables)))
sample = tf.nn.sigmoid(sample_logit)
output = {'predicted_occ': sample, 'predicted_free': 1 - sample}
metrics = nn.calc_metrics(output, batch)
metrics['loss/angle'] = angle_loss
metrics['loss/vae'] = vae_loss
metrics['values/variance'] = tf.reduce_mean(tf.exp(logvar_angle))
m = {k: tf.reduce_mean(metrics[k]) for k in metrics}
m['loss'] = loss
return output, m
def step_fn(batch):
with tf.GradientTape() as tape:
output_logits = self.model(self.prep_input(batch))
sample = tf.nn.sigmoid(output_logits)
output = {
'predicted_occ': sample,
'predicted_free': 1 - sample
}
metrics = nn.calc_metrics(output, batch)
cross_ent = tf.nn.sigmoid_cross_entropy_with_logits(
logits=output_logits, labels=batch['gt_occ'])
loss = nn.reduce_sum_batch(cross_ent)
variables = self.model.trainable_variables
gradients = tape.gradient(loss, variables)
clipped_gradients = [
tf.clip_by_value(g, -1e6, 1e6) for g in gradients
]
self.opt.apply_gradients(
list(zip(clipped_gradients, variables)))
return loss, metrics
def step_fn_multiloss(batch):
with tf.GradientTape() as tape:
output = self(batch)
metrics = nn.calc_metrics(output, batch)
loss = tf.reduce_sum(
tf.keras.losses.binary_crossentropy(
batch['gt_occ'], output['predicted_occ']))
loss = loss / self.batch_size
metrics['loss/0_step'] = loss
m = metrics
# for i in range(1):
# Multistep part
if True:
i = 0
b = {'conditioned_occ': output['predicted_occ']}
output = self(b)
step_loss = tf.reduce_sum(
tf.keras.losses.binary_crossentropy(
batch['gt_occ'], output['predicted_occ']))
step_loss = step_loss / self.batch_size
metrics['loss/{}_step'.format(i + 1)] = step_loss
loss = tf.cond(
metrics['pred|gt/p(predicted_occ|gt_occ)'] > 0.95,
lambda: tf.add(step_loss, loss), lambda: loss)
variables = self.model.trainable_variables
gradients = tape.gradient(loss, variables)
clipped_gradients = [
tf.clip_by_value(g, -1e6, 1e6) for g in gradients
]
self.opt.apply_gradients(
list(zip(clipped_gradients, variables)))
return loss, metrics
def step_fn(batch):
with tf.GradientTape() as tape:
known = stack_known(batch)
mean, logvar = self.encode(known)
z = self.reparameterize(mean, logvar)
sample_logit = self.decode(z)
vae_loss = compute_vae_loss(z,
mean,
logvar,
sample_logit,
labels=batch['gt_occ'])
sample = tf.nn.sigmoid(sample_logit)
output = {
'predicted_occ': sample,
'predicted_free': 1 - sample
}
metrics = nn.calc_metrics(output, batch)
vae_variables = self.encoder.trainable_variables + self.generator.trainable_variables
gradients = tape.gradient(vae_loss, vae_variables)
self.opt.apply_gradients(list(zip(gradients, vae_variables)))
return vae_loss, metrics