def pretrain(self, optimizer, lr, run_optimizer=True):
"""Constructs training ops for pre-training the Generator.
RNN-RBM allows pre-training the RBM module.
Args:
optimizer: tf.Optimizer object that computes gradients.
lr: A learning rate for weight updates.
run_optimizer: `bool` indicating whether to run optimizer
on the generator's trainable variables, or return
an empty update ops list. Allows running the optimizer
on the top of multiple generators.
Returns:
init_ops: A list of weight initialization ops.
update_ops: A list of weights update ops.
metrics: A dictionary of metric names and metric ops.
metrics_upd: A list of metric update ops.
summaries: A dict of tf.Summary objects for model's
metrics, weights and training gradients.
"""
inputs_flat = flatten_maybe_padded_sequences(self._inputs,
self._lengths)
return self._rbm.train(inputs_flat, lr)
def _get_state(self,
inputs,
lengths=None,
initial_state=None,
last_outputs=False):
"""Computes the new state of the RNN-NADE.
Scans through the inputs and returns the final state that either
considers only the last outputs, or the outputs for all sequence steps.
The final state consists of Estimator bias parameters and RNN cell's state.
Args:
inputs: A batch of sequences to process,
sized `[batch_size, time_steps, num_input_dims]`
or `[batch_size, num_input_dims]`.
lengths: The lengths of input sequences sized `[batch_size]`,
or None if all are equal.
initial_state: An RnnEstimatorStateTuple, the initial state of the RNN-NADE,
or None if the zero state should be used.
last_outputs: `bool` indicating whether to return the outputs
only from the last time step, or from all sequence steps.
Returns:
final_state: An RnnEstimatorStateTuple, the final state of the RNN-NADE.
"""
# Processing the `lengths` and `initial_state` arguments
lengths, initial_rnn_state = self._get_state_helper(
inputs, lengths, initial_state)
helper = tf.contrib.seq2seq.TrainingHelper(inputs=inputs,
sequence_length=lengths)
decoder = tf.contrib.seq2seq.BasicDecoder(
cell=self._rnn_cell,
helper=helper,
initial_state=initial_rnn_state,
output_layer=self._fc_layer)
# Scanning through the input sequences
with tf.variable_scope(get_current_scope()):
final_outputs, final_rnn_state = tf.contrib.seq2seq.dynamic_decode(
decoder)[0:2]
# Final flattened outputs from the RNN-Dense decoder
with tf.variable_scope('final_outputs'):
if last_outputs:
final_outputs_flat = final_outputs.rnn_output[:, -1, :]
else:
final_outputs_flat = flatten_maybe_padded_sequences(
final_outputs.rnn_output, lengths)
# Compute NADE biases from the Dense layer outputs.
with tf.variable_scope('nade_biases'):
b_enc, b_dec = self._build_biases(final_outputs_flat)
with tf.variable_scope('final_state'):
return RnnEstimatorStateTuple(b_enc, b_dec, final_rnn_state)
def _build_targets(self):
"""Builds prediction targets from inputs.
Returns:
inputs: A batch of stacked track target predictions.
"""
return tf.reshape(
flatten_maybe_padded_sequences(self._x, self._lengths),
[-1, self.num_dims * self.num_tracks])
def _get_state(self,
inputs,
lengths=None,
initial_state=None,
last_outputs=False):
"""Computes the new state of the RNN-RBM.
Scans through the inputs and returns the final state that either
considers only the last outputs, or the outputs for all sequence steps.
The final state consists of Estimator bias parameters and RNN cell's state.
Args:
inputs: A batch of sequences to process,
sized `[batch_size, time_steps, num_input_dims]`
or `[batch_size, num_input_dims]`.
lengths: The lengths of input sequences sized `[batch_size]`,
or None if all are equal.
initial_state: An RnnEstimatorStateTuple, the initial state of the RNN-RBM,
or None if the zero state should be used.
last_outputs: `bool` indicating whether to return the outputs
only from the last time step, or from all sequence steps.
Returns:
final_state: An RnnEstimatorStateTuple, the final state of the RNN-RBM.
"""
# Processing the `lengths` and `initial_state` arguments
lengths, initial_rnn_state = self._get_state_helper(
inputs, lengths, initial_state)
# Scanning through the input sequences
with tf.variable_scope(f'{get_current_scope()}{self.track_name}'):
final_outputs, final_rnn_state = tf.nn.dynamic_rnn(
self._rnn_cell,
inputs,
initial_state=initial_rnn_state,
sequence_length=lengths,
parallel_iterations=32,
dtype=tf.float32)
# Final flattened outputs from the RNN-Dense decoder
with tf.variable_scope('final_outputs'):
if last_outputs:
final_outputs = final_outputs[:, -1, :]
final_outputs_flat = tf.reshape(
final_outputs, [-1, tf.shape(final_outputs)[-1]])
else:
final_outputs_flat = flatten_maybe_padded_sequences(
final_outputs, lengths)
# Compute RBM biases from the RNN cell outputs.
with tf.variable_scope('rbm_biases'):
bh_t, bv_t = self._build_biases(final_outputs_flat)
with tf.variable_scope('final_state'):
return RnnEstimatorStateTuple(bh_t, bv_t, final_rnn_state)
def _build_targets(self):
"""Builds prediction targets from inputs.
Returns:
inputs: A list of per-track target predictions.
"""
targets_flat = flatten_maybe_padded_sequences(self._x, self._lengths)
# Transforming into a list of per-track predictions
return tf.unstack(targets_flat, axis=-1)
def train(self, optimizer, lr, layer=0):
"""Constructs training ops for the DBNEncoder.
Args:
optimizer: tf.Optimizer object that computes gradients.
lr: A learning rate for weight updates.
layer: A layer index for which the parameters should be updated.
Returns:
init_ops: A list of weight initialization ops.
update_ops: A list of weights update ops.
metrics: A dictionary of metric names and metric ops.
metrics_upd: A list of metric update ops.
summaries: A dict of tf.Summary objects for model's
metrics, weights and training gradients.
"""
assert 0 <= layer < self.num_layers
# RBM to be trained
rbm = self.dbn.rbm_layers[layer]
# Accessing inputs to the `layer`-th RBM.
x = self._inputs if layer == 0 else self.encodings[layer - 1]
x = flatten_maybe_padded_sequences(x, self._lengths)
# Building hidden vectors and reconstructions for the given RBM.
with tf.variable_scope(f'{rbm.name}-{layer}/{rbm.track_name}'):
with tf.variable_scope('encodings'):
_, encodings = rbm.forward(x)
with tf.variable_scope('decodings'):
dec_probs, decodings, = rbm.reconstruct(encodings)
# Building metrics and summaries
metrics, metrics_upd, metrics_summaries = rbm.build_metrics(
targets=x, predictions=decodings, cond_probs=dec_probs)
summaries = {
'weights': rbm.weight_summary,
'metrics': metrics_summaries,
}
# Getting update ops
init_ops, update_ops, summaries['gradients'] = rbm.train(x, lr)
return init_ops, update_ops, metrics, metrics_upd, summaries
def build(self, x=None, y=None, lengths=None, is_train=None, mode='eval'):
"""Builds RNN-MultiNADE's graph.
Defines the forward pass through the input sequences.
Note:
The inputs to the model are zero padded for the first time step.
Args:
x: A batch of input sequences to the generator,
sized `[batch_size, time_steps, ...]`.
y: A batch of target predictions for the generator,
sized `[batch_size, time_steps, num_dims]`.
lengths: The lengths of input sequences sized `[batch_size]`.
is_train: `bool` for distinguishing between training and inference.
mode: Build mode for optimizing the graph size.
Some operations may not be needed for training or inference.
Raises:
ValueError: If incorrect build mode is provided.
"""
RnnEstimator.build(self, x, y, lengths, is_train, mode)
if mode in ('train', 'eval'):
inputs, targets = x, y
# Flatten targets into the shape `[sum(lengths), num_dims]`
# and split into a list of track targets
with tf.variable_scope(f'{get_current_scope()}generator_targets/flat/'):
targets_flat = tf.reshape(
flatten_maybe_padded_sequences(targets, lengths),
[-1, self.num_dims, self.num_tracks]
)
targets_flat = tf.unstack(targets_flat, axis=-1)
# Compute log probabilities for inputs and targets by scanning through inputs
with tf.name_scope(f'{get_current_scope()}{self.name}/'):
log_probs, cond_probs = self.log_prob(inputs, targets_flat, lengths)
# Model outputs binarization
with tf.variable_scope(f'{get_current_scope()}generator_outputs/'):
self._outputs = []
for i in range(self.num_tracks):
with tf.variable_scope(self.tracks[i]):
self._outputs.append(tf.cast(tf.greater_equal(cond_probs[i], .5), tf.float32))
# Building model metrics and summaries from the targets and outputs
self._metrics, self._metrics_upd, self._summaries = self.build_metrics(
targets=targets_flat,
predictions=self._outputs,
cond_probs=cond_probs,
log_probs=log_probs
)
# RNN-MultiNADE variables
self._variables = {
'rnn': self._rnn.variables,
'nade': [nade.variables for nade in self._nades],
'dense': self._fc_layer.trainable_variables
}
self._trainable_variables = \
self._rnn.trainable_variables \
+ list(itertools.chain.from_iterable([nade.trainable_variables for nade in self._nades])) \
+ self._fc_layer.trainable_variables
self._summaries['weights'] = self.weight_summary
self._is_built = True
def build(self, x=None, y=None, lengths=None, is_train=None, mode='eval'):
"""Builds RNN-RBM's graph.
Defines the forward pass through the input sequences.
Note:
The inputs to the model are zero padded for the first time step.
Args:
x: A batch of input sequences to the generator,
sized `[batch_size, time_steps, ...]`.
y: A batch of target predictions for the generator,
sized `[batch_size, time_steps, num_dims]`.
lengths: The lengths of input sequences sized `[batch_size]`.
is_train: `bool` for distinguishing between training and inference.
mode: Build mode for optimizing the graph size.
Some operations may not be needed for training or inference.
Raises:
ValueError: If incorrect build mode is provided.
"""
super().build(x, y, lengths, is_train, mode)
if mode in ('train', 'eval'):
inputs, targets = x, y
# Flatten inputs into the shape `[sum(lengths), num_dims]`
with tf.variable_scope(
f'{get_current_scope()}generator_inputs/{self.track_name}/flat/'
):
inputs_flat = flatten_maybe_padded_sequences(inputs, lengths)
# Flatten targets into the shape `[sum(lengths), num_dims]`
with tf.variable_scope(
f'{get_current_scope()}generator_targets/{self.track_name}/flat/'
):
targets_flat = flatten_maybe_padded_sequences(targets, lengths)
# Scanning through the input sequences.
with tf.name_scope(f'{get_current_scope()}{self.name}/'):
initial_state = self.zero_state(batch_size=tf.shape(x)[0])
final_state = self._get_state(inputs,
initial_state=initial_state)
# Sampling outputs from the RBM
with tf.variable_scope(
f'{get_current_scope()}generator_outputs/{self.track_name}/'
):
self._outputs, cond_probs = self.sample_single(
inputs_flat, final_state)
# Building model metrics and summaries from the targets and outputs
self._metrics, self._metrics_upd, self._summaries[
'metrics'] = self.build_metrics(
targets=targets_flat,
predictions=self._outputs,
cond_probs=cond_probs,
)
# RNN-RBM variables
self._variables = {
**self._rbm.variables,
'Wuh': self._Wuh,
'Wuv': self._Wuv,
**self._rnn.variables,
}
self._trainable_variables = \
self._rbm.trainable_variables \
+ self._rnn.trainable_variables \
+ [self._Wuh, self._Wuv]
self._summaries['weights'] = self.weight_summary
self._is_built = True