def init_weights(self):
""" Initialize all the trainable weights."""
self.w_g0 = weight_variable((self.sensor_size, self.hg_size))
self.b_g0 = bias_variable((self.hg_size, ))
self.w_g1 = weight_variable((self.hg_size, self.g_size))
self.b_g1 = bias_variable((self.g_size, ))
self.w_l0 = weight_variable((self.loc_dim, self.g_size))
self.b_l0 = bias_variable((self.g_size, ))
def init_weights(self):
self.w = weight_variable((self.input_dim, self.loc_dim))
self.b = bias_variable((self.loc_dim, ))
def buildModel(self):
inputShape = self.params.original_size
#Running on GPU
with tf.device(self.params.device):
with tf.name_scope("Input"):
self.images = tf.placeholder(tf.float32,
shape=[
None, inputShape[0] *
inputShape[1] * inputShape[2]
],
name="images")
reshape_images = tf.reshape(
self.images,
[-1, inputShape[0], inputShape[1], inputShape[2]])
tf.summary.image('images', reshape_images)
self.varDict['images'] = self.images
self.labels = tf.placeholder(tf.int64,
shape=[None],
name="labels")
self.varDict['labels'] = self.labels
#Build core network
with tf.variable_scope('fc_net'):
w_fc1 = weight_variable(
(inputShape[0] * inputShape[1] * inputShape[2],
self.params.num_fc1_units))
b_fc1 = bias_variable((self.params.num_fc1_units, ))
fc1 = tf.nn.relu(tf.matmul(self.images, w_fc1) + b_fc1)
w_fc2 = weight_variable(
(self.params.num_fc1_units, self.params.num_fc2_units))
b_fc2 = bias_variable((self.params.num_fc2_units, ))
fc2 = tf.nn.relu(tf.matmul(fc1, w_fc2) + b_fc2)
#Classification network
with tf.variable_scope('classification_net'):
output = fc2
#Pass through classification network for reward
w_logit = weight_variable(
(self.params.num_fc2_units, self.params.num_classes))
b_logit = bias_variable((self.params.num_classes, ))
logits = tf.nn.xw_plus_b(output, w_logit, b_logit)
self.softmax = tf.nn.softmax(logits)
with tf.variable_scope('loss'):
# cross-entropy.
xent = tf.nn.sparse_softmax_cross_entropy_with_logits(
logits=logits, labels=self.labels)
xent = tf.reduce_mean(xent)
pred_labels = tf.argmax(logits, 1)
#hybrid loss
loss = xent
var_list = tf.trainable_variables()
grads = tf.gradients(loss, var_list)
grads, _ = tf.clip_by_global_norm(grads,
self.params.max_grad_norm)
#Add to scalar tensorboard
self.scalarDict['xent'] = xent
self.scalarDict['loss'] = loss
with tf.variable_scope('accuracy'):
#
self.injectBool = tf.placeholder_with_default(
False, shape=(), name="injectBool")
self.injectAcc = tf.placeholder_with_default(0.0,
shape=None,
name="injectAcc")
#Calculate accuracy
pred_labels = tf.argmax(self.softmax, axis=1)
calc_acc = tf.reduce_mean(
tf.cast(tf.equal(pred_labels, self.labels), tf.float32))
accuracy = tf.cond(self.injectBool, lambda: self.injectAcc,
lambda: calc_acc)
self.scalarDict['accuracy'] = accuracy
with tf.variable_scope('optimizer'):
# learning rate
global_step = tf.get_variable(
'global_step', [],
initializer=tf.constant_initializer(0),
trainable=False)
training_steps_per_epoch = self.params.num_train_examples // self.params.batch_size
# decay per training epoch
learning_rate = tf.train.exponential_decay(
self.params.lr_start,
global_step,
training_steps_per_epoch,
self.params.lr_decay,
staircase=True)
learning_rate = tf.maximum(learning_rate, self.params.lr_min)
opt = tf.train.AdamOptimizer(learning_rate)
self.train_op = opt.apply_gradients(zip(grads, var_list),
global_step=global_step)
self.scalarDict['learning_rate'] = learning_rate
def buildModel(self):
inputShape = self.params.original_size
#Running on GPU
with tf.device(self.params.device):
with tf.name_scope("Input"):
self.images = tf.placeholder(tf.float32,
shape=[
None, inputShape[0] *
inputShape[1] * inputShape[2]
],
name="images")
reshape_images = tf.reshape(
self.images,
[-1, inputShape[0], inputShape[1], inputShape[2]])
#self.addImageSummary('images', reshape_images, False)
tf.summary.image('images', reshape_images)
self.varDict['images'] = self.images
self.rescale_images = self.images * 2 - 1
self.labels = tf.placeholder(tf.int64,
shape=[None],
name="labels")
self.varDict['labels'] = self.labels
self.is_train = tf.placeholder(tf.bool,
shape=(),
name="is_train")
#Build aux nets
# Build the aux nets.
with tf.variable_scope('glimpse_net'):
self.gl = GlimpseNet(self.params, self.rescale_images)
#Add variables to varDict
self.varDict.update(self.gl.getVars())
with tf.variable_scope('loc_net'):
self.loc_net = LocNet(self.params)
#Add variables to varDict
self.varDict.update(self.loc_net.getVars())
#Build core network
with tf.variable_scope('core_net'):
# number of examples
N = tf.shape(self.images)[0]
#Initial glimpse location
#TODO this should be calculated from context network
init_loc = tf.random_uniform((N, 2), minval=-1, maxval=1)
#init_loc = tf.zeros((N, 2))
#RNN base building block
lstm_cell_0 = tf.contrib.rnn.BasicLSTMCell(
self.params.cell_size_0)
lstm_cell_1 = tf.contrib.rnn.BasicLSTMCell(
self.params.cell_size_1)
#Initial state
state_0 = lstm_cell_0.zero_state(N, tf.float32)
#TODO this should be initialized from context network
state_1 = lstm_cell_1.zero_state(N, tf.float32)
current_glimpse = self.gl(init_loc, self.is_train)
outputs_0 = []
outputs_1 = []
glimpses = [current_glimpse]
for i in range(self.params.num_glimpses):
output_0, state_0 = lstm_cell_0(current_glimpse, state_0)
#output_0 = batchnorm(output_0, "core_0", self.is_train)
outputs_0.append(output_0)
output_1, state_1 = lstm_cell_1(output_0, state_1)
#output_1 = batchnorm(output_1, "core_1", self.is_train)
outputs_1.append(output_1)
#Update glimpse
current_glimpse = self.get_next_input(output_1, i)
glimpses.append(current_glimpse)
self.varDict["core_output_0"] = outputs_0
self.varDict["core_output_1"] = outputs_1
self.varDict["glimpses"] = glimpses
#Baseline for variance reduction
with tf.variable_scope('baseline'):
baseline_w = weight_variable((self.params.cell_size_1, 1))
baseline_b = bias_variable((1, ))
baselines = []
for t, output in enumerate(outputs_1):
baseline_t = tf.nn.xw_plus_b(output, baseline_w,
baseline_b)
baseline_t = tf.squeeze(baseline_t)
baselines.append(baseline_t)
baselines = tf.stack(baselines) # [timesteps, batch_sz]
baselines = tf.transpose(baselines) # [batch_sz, timesteps]
self.varDict["baseline_w"] = baseline_w
self.varDict["baseline_b"] = baseline_b
self.varDict["baseline_h"] = baselines
#Classification network
with tf.variable_scope('classification_net'):
# Take the last step of lower rnn only.
output = outputs_0[-1]
#Pass through classification network for reward
classnet_w = weight_variable(
(self.params.cell_size_0, self.params.classnet_size))
classnet_b = bias_variable((self.params.classnet_size, ))
logit_w = weight_variable(
(self.params.classnet_size, self.params.num_classes))
logit_b = bias_variable((self.params.num_classes, ))
classnet_h = tf.nn.relu(
tf.nn.xw_plus_b(output, classnet_w, classnet_b))
logits = tf.nn.xw_plus_b(classnet_h, logit_w, logit_b)
self.softmax = tf.nn.softmax(logits)
self.varDict["classnet_w"] = classnet_w
self.varDict["classnet_b"] = classnet_b
self.varDict["classnet_h"] = classnet_h
self.varDict["logit_w"] = logit_w
self.varDict["logit_b"] = logit_b
self.varDict["logits"] = logits
with tf.variable_scope('loss'):
# cross-entropy.
xent = tf.nn.sparse_softmax_cross_entropy_with_logits(
logits=logits, labels=self.labels)
xent = tf.reduce_mean(xent)
pred_labels = tf.argmax(logits, 1)
# 0/1 reward.
reward = tf.cast(tf.equal(pred_labels, self.labels),
tf.float32)
rewards = tf.expand_dims(reward, 1) # [batch_sz, 1]
rewards = tf.tile(
rewards,
(1, self.params.num_glimpses)) # [batch_sz, timesteps]
logll = loglikelihood(self.loc_mean_arr, self.sampled_loc_arr,
self.params.loc_std)
#Only update rewards with reinforce
advs = rewards - tf.stop_gradient(baselines)
logllratio = tf.reduce_mean(logll * advs)
reward = tf.reduce_mean(reward)
#Baseline loss
#Only update baselines with baseline_mse
baselines_mse = tf.reduce_mean(
tf.square((tf.stop_gradient(rewards) - baselines)))
var_list = tf.trainable_variables()
#hybrid loss
#TODO baselines only update based on mse?
hybrid_loss = -self.params.reinforce_lambda * logllratio + baselines_mse + xent
grads = tf.gradients(hybrid_loss, var_list)
grads, _ = tf.clip_by_global_norm(grads,
self.params.max_grad_norm)
#Add to scalar tensorboard
self.scalarDict['hybrid_loss'] = hybrid_loss
self.scalarDict['xent'] = xent
self.scalarDict['baseline_mse'] = baselines_mse
self.scalarDict['reinforce_loss'] = -logllratio
self.scalarDict['reward'] = reward
self.varDict['logll'] = logll
self.varDict['loc_mean_arr'] = self.loc_mean_arr
self.varDict['sampled_loc_arr'] = self.sampled_loc_arr
with tf.variable_scope('accuracy'):
#
self.injectBool = tf.placeholder_with_default(
False, shape=(), name="injectBool")
self.injectAcc = tf.placeholder_with_default(0.0,
shape=None,
name="injectAcc")
#Calculate accuracy
pred_labels = tf.argmax(self.softmax, axis=1)
calc_acc = tf.reduce_mean(
tf.cast(tf.equal(pred_labels, self.labels), tf.float32))
accuracy = tf.cond(self.injectBool, lambda: self.injectAcc,
lambda: calc_acc)
self.scalarDict['accuracy'] = accuracy
with tf.variable_scope('optimizer'):
# learning rate
global_step = tf.get_variable(
'global_step', [],
initializer=tf.constant_initializer(0),
trainable=False)
training_steps_per_epoch = self.params.num_train_examples // self.params.batch_size
# decay per training epoch
learning_rate = tf.train.exponential_decay(
self.params.lr_start,
global_step,
training_steps_per_epoch,
self.params.lr_decay,
staircase=True)
learning_rate = tf.maximum(learning_rate, self.params.lr_min)
opt = tf.train.AdamOptimizer(learning_rate)
#opt = tf.train.MomentumOptimizer(learning_rate, .9, use_nesterov=True)
#Makes sure to update batchnorm var before training step
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops):
self.train_op = opt.apply_gradients(
zip(grads, var_list), global_step=global_step)
self.scalarDict['learning_rate'] = learning_rate