def build_tower(self, build_backprop):
with tf.variable_scope("Inputs"):
# Input placeholders [Batch, Length, X]
placeholders = {}
placeholders["sequences"] = \
tf.placeholder(tf.float32, [None, None, self.dims["alphabet"]],
name="Sequences")
placeholders["secondary_structure"] = \
tf.placeholder(tf.float32, [None, None, self.dims["ss"]],
name="SecondaryStructure")
placeholders["coordinates_target"] = \
tf.placeholder(tf.float32, [None, None, self.dims["coords"]],
name="TargetCoordinates")
placeholders["lengths"] = \
tf.placeholder(tf.int32, [None], name="SequenceLengths")
placeholders["structure_mask"] = \
tf.placeholder(tf.float32, [None, None], name="StructureMask")
placeholders["sequences"] = tf.identity(placeholders["sequences"])
placeholders["bptt_log_decay"] = self.bptt_log_decay
if self.use_ec:
placeholders["couplings"] = tf.placeholder(tf.float32,
[None, None, None],
name="Couplings")
# Optional placeholders
default_dropout = tf.constant(
self.hyperparams["training"]["dropout"], dtype=tf.float32)
placeholders["global_step"] = self.global_step
placeholders["training"] = tf.placeholder_with_default(
False, (), name="Training")
placeholders["dropout"] = tf.placeholder_with_default(
default_dropout, (), name="Dropout")
placeholders["native_init_prob"] = tf.placeholder_with_default(
0., (), name="NativeInitProb")
placeholders["native_unfold_max"] = tf.placeholder_with_default(
0.5, (), name="NativeUnfoldMax")
placeholders[
"native_unfold_randomize"] = tf.placeholder_with_default(
False, (), name="NativeRandomUnfold")
placeholders["langevin_steps"] = tf.placeholder_with_default(
self.hyperparams["folding"]["langevin_steps"], (),
name="LangevinSteps")
placeholders["beta_anneal"] = tf.placeholder_with_default(
1.0, (), name="BetaAnneal")
# Update unknown dims
dims = self.dims
seq_shape = tf.shape(placeholders["sequences"])
dims["batch"] = seq_shape[0]
dims["length"] = seq_shape[1]
# Keep track of which variables we create
prior_vars = tf.trainable_variables()
# Build the protein folding engine
physics = Physics(placeholders, dims, self.global_step,
self.hyperparams)
physics.build_graph()
tensors = physics.tensors
# Build the loss function
loss = Loss(placeholders, tensors, dims, self.global_step,
self.hyperparams)
loss.build_graph()
tensors["score"] = loss.tensors["score"]
tensors["coarse_target"] = loss.tensors["coarse_target"]
new_vars = tf.trainable_variables()
params = {}
params["score"] = [v for v in new_vars if "Score" in v.name]
params["generator"] = [v for v in new_vars if v not in params["score"]]
# Parameter accounting
for (param_name, param_set) in params.iteritems():
print "Parameter set: " + param_name
p_counts = [np.prod(v.get_shape().as_list()) for v in param_set]
p_names = [v.name for v in param_set]
p_total = sum(p_counts)
print " contains " + str(p_total) + " params in " \
+ str(len(p_names)) + " tensors"
print "\n"
# Backpropagate
gradients = None
train_ops = physics.layers.train_ops + loss.layers.train_ops
if build_backprop:
with tf.variable_scope("Backprop"):
with tf.variable_scope("Generator"):
g_loss = tensors["score"]
g_vars = params["generator"] + params["score"]
g_grads = tf.gradients(g_loss, g_vars)
gen_gvs = zip(g_grads, g_vars)
gradients, gen_grad_norm = self.clip_gradients(gen_gvs)
""" Gradient damping (gamma) adaptation for simulator
During simulator roll-outs, backpropgation is damped as
gamma = exp(log_decay)
X = gamma * X + (1. - gamma) * tf.stop_gradient(X)
When gamma=1.0, this behaves in the usual manner.
[ Current Heuristic ]
When the model is 'stable':
slowly raise gamma with log-geometric scaling
When the model is 'unstable':
rapidly lower gamma with log-linear scaling
"""
is_stable = tf.logical_and(tf.is_finite(gen_grad_norm),
gen_grad_norm < 100.)
decay_update = tf.cond(is_stable,
lambda: 0.99 * self.bptt_log_decay,
lambda: self.bptt_log_decay - 0.01)
self.bptt_decay_ops += [
self.bptt_log_decay.assign(decay_update)
]
tf.summary.scalar("GradientNorm",
gen_grad_norm,
collections=["gen_batch", "gen_dense"])
tf.summary.scalar("BPTTLogDecay",
self.bptt_log_decay,
collections=["gen_batch", "gen_dense"])
return placeholders, tensors, gradients, train_ops
