def __build_model_and_summary(self):
# split the tensors
with tf.variable_scope("tower_split"), tf.device("/cpu:0"):
# splitted tensors
ts_inputs = tf.split(self.t_inputs, self.num_gpu, 1)
ts_lstm_initial_state = tf.split(self.t_lstm_initial_state, self.num_gpu, 2)
ts_initial_poses = tf.split(self.t_initial_poses, self.num_gpu, 0)
ts_imu_data = tf.split(self.t_imu_data, self.num_gpu, 1)
ts_ekf_initial_state = tf.split(self.t_ekf_initial_state, self.num_gpu, 0)
ts_ekf_initial_covar = tf.split(self.t_ekf_initial_covariance, self.num_gpu, 0)
ts_se3_labels = tf.split(self.t_se3_labels, self.num_gpu, 1)
ts_fc_labels = tf.split(self.t_fc_labels, self.num_gpu, 1)
# list to store results
ts_ekf_states = []
ts_ekf_covar_states = []
ts_lstm_states = []
losses_keys = ["se3_loss", "se3_xyz_loss", "se3_quat_loss",
"fc_loss", "fc_xyz_loss", "fc_ypr_loss", "x_loss", "y_loss", "z_loss",
"total_loss"]
ts_losses_dict = dict(zip(losses_keys, [[] for i in range(len(losses_keys))]))
for i in range(0, self.num_gpu):
device_setter = tf.train.replica_device_setter(ps_tasks=1,
ps_device='/job:localhost/replica:0/task:0/device:CPU:0',
worker_device='/job:localhost/replica:0/task:0/device:GPU:%d' % i)
with tf.name_scope("tower_%d" % i), tf.device(device_setter):
tools.printf("Building model...")
fc_outputs, fc_covar, se3_outputs, lstm_states, ekf_states, ekf_covar_states = \
model.build_seq_model(self.cfg, ts_inputs[i], ts_lstm_initial_state[i], ts_initial_poses[i],
ts_imu_data[i], ts_ekf_initial_state[i], ts_ekf_initial_covar[i],
self.t_is_training, get_activations=True,
use_initializer=self.t_use_initializer,
use_ekf=self.cfg.use_ekf)
# this returns lstm states as a tuple, we need to stack them
lstm_states = tf.stack(lstm_states, 0)
ts_lstm_states.append(lstm_states)
ts_ekf_states.append(ekf_states)
ts_ekf_covar_states.append(ekf_covar_states)
with tf.variable_scope("loss"):
se3_loss, se3_xyz_loss, se3_quat_loss \
= losses.se3_losses(se3_outputs, ts_se3_labels[i], self.cfg.k_se3)
fc_loss, fc_xyz_loss, fc_ypr_loss, x_loss, y_loss, z_loss \
= losses.fc_losses(fc_outputs, fc_covar, ts_fc_labels[i], self.cfg.k_fc)
total_loss = (1 - self.t_alpha) * se3_loss + self.t_alpha * fc_loss
for k, v in ts_losses_dict.items():
v.append(locals()[k])
tf.get_variable_scope().reuse_variables()
with tf.variable_scope("tower_join"), tf.device("/cpu:0"):
# join the lstm states
self.t_lstm_states = tf.concat(ts_lstm_states, 2)
for k, v in ts_losses_dict.items():
ts_losses_dict[k] = tf.reduce_mean(v)
self.t_ekf_states = tf.concat(ts_ekf_states, 0)
self.t_ekf_covar_states = tf.concat(ts_ekf_covar_states, 0)
self.t_total_loss = ts_losses_dict["total_loss"]
self.t_se3_loss = ts_losses_dict["se3_loss"]
tools.printf("Building optimizer...")
with tf.variable_scope("optimizer", reuse=tf.AUTO_REUSE):
if self.cfg.use_init and self.cfg.only_train_init:
train_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, "initializer_layer")
elif self.cfg.train_noise_covariance and self.cfg.static_nn:
train_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, "imu_noise_params")
else:
train_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES)
self.op_trainer = tf.train.AdamOptimizer(learning_rate=self.t_lr). \
minimize(self.t_total_loss, colocate_gradients_with_ops=True, var_list=train_vars)
# tensorboard summaries
tools.printf("Building tensorboard summaries...")
with tf.device("/cpu:0"):
self.t_sequence_id = tf.placeholder(dtype=tf.uint8, shape=[])
self.t_epoch = tf.placeholder(dtype=tf.int32, shape=[])
tf.summary.scalar("total_loss", ts_losses_dict["total_loss"])
tf.summary.scalar("fc_loss", ts_losses_dict["fc_loss"])
tf.summary.scalar("se3_loss", ts_losses_dict["se3_loss"])
tf.summary.scalar("fc_xyz_loss", ts_losses_dict["fc_xyz_loss"])
tf.summary.scalar("fc_ypr_loss", ts_losses_dict["fc_ypr_loss"])
tf.summary.scalar("se3_xyz_loss", ts_losses_dict["se3_xyz_loss"])
tf.summary.scalar("se3_quat_loss", ts_losses_dict["se3_quat_loss"])
tf.summary.scalar("x_loss", ts_losses_dict["x_loss"])
tf.summary.scalar("y_loss", ts_losses_dict["y_loss"])
tf.summary.scalar("z_loss", ts_losses_dict["z_loss"])
tf.summary.scalar("alpha", self.t_alpha)
tf.summary.scalar("lr", self.t_lr)
tf.summary.scalar("sequence_id", self.t_sequence_id)
tf.summary.scalar("epoch", self.t_epoch)
self.op_train_merged_summary = tf.summary.merge_all()
activations = tf.get_collection(tf.GraphKeys.ACTIVATIONS)
initial_layer = tf.summary.image("1st layer activations", tf.expand_dims(activations[0][:, 0, :, :], -1))
final_layer = tf.summary.image("Last layer activations", tf.expand_dims(activations[1][:, 0, :, :], -1))
self.op_train_image_summary = tf.summary.merge([initial_layer, final_layer])
val_loss_sum = tf.summary.scalar("val_total_loss", ts_losses_dict["total_loss"])
val_fc_sum = tf.summary.scalar("val_fc_losses", ts_losses_dict["fc_loss"])
val_se3_sum = tf.summary.scalar("val_se3_losses", ts_losses_dict["se3_loss"])
val_fc_xyz_loss = tf.summary.scalar("val_fc_xyz_loss", ts_losses_dict["fc_xyz_loss"])
val_fc_ypr_loss = tf.summary.scalar("val_fc_ypr_loss", ts_losses_dict["fc_ypr_loss"])
val_se3_xyz_loss = tf.summary.scalar("val_se3_xyz_loss", ts_losses_dict["se3_xyz_loss"])
val_se3_quat_loss = tf.summary.scalar("val_se3_quat_loss", ts_losses_dict["se3_quat_loss"])
val_x_sum = tf.summary.scalar("val_x_loss", ts_losses_dict["x_loss"])
val_y_sum = tf.summary.scalar("val_y_loss", ts_losses_dict["y_loss"])
val_z_sum = tf.summary.scalar("val_z_loss", ts_losses_dict["z_loss"])
self.op_val_merged_summary = tf.summary.merge(
[val_loss_sum, val_fc_sum, val_se3_sum, val_fc_xyz_loss, val_fc_ypr_loss, val_se3_xyz_loss,
val_se3_quat_loss, val_x_sum, val_y_sum, val_z_sum])
fc_labels = tf.placeholder(tf.float32,
name="se3_labels",
shape=[timesteps, batch_size, 6])
# dynamic learning rates
se3_lr = tf.placeholder(tf.float32, name="se3_lr", shape=[])
fc_lr = tf.placeholder(tf.float32, name="fc_lr", shape=[])
# =================== MODEL + LOSSES + Optimizer ========================
fc_outputs, se3_outputs, lstm_states = model.build_model(
inputs, lstm_init_state)
with tf.device("/gpu:0"):
with tf.variable_scope("Losses"):
se3_losses = losses.se3_losses(se3_outputs, se3_labels, k)
fc_losses = losses.fc_losses(fc_outputs, fc_labels)
with tf.variable_scope("Optimizer"):
with tf.device("/gpu:0"):
se3_trainer = tf.train.AdamOptimizer(
learning_rate=se3_lr).minimize(se3_outputs)
with tf.device("/gpu:0"):
fc_trainer = tf.train.AdamOptimizer(
learning_rate=fc_lr).minimize(fc_outputs)
# =================== TRAINING ========================
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
# Visualization
writer = tf.summary.FileWriter('graph_viz/')