def __init__(self, arch, is_training=False):
'''
Variational auto-encoder implemented in 2D convolutional neural nets
Input:
`arch`: network architecture (`dict`)
`is_training`: (unused now) it was kept for historical reasons (for `BatchNorm`)
'''
self.arch = arch
self._sanity_check()
self.is_training = is_training
with tf.name_scope('SpeakerRepr'):
self.y_emb = self._l2_regularized_embedding(
self.arch['y_dim'],
self.arch['z_dim'],
'y_embedding')
self._generate = tf.make_template(
'Generator',
self._generator)
self._encode = tf.make_template(
'Encoder',
self._encoder)
self.generate = self.decode # for VAE-GAN extension
def _build_networks(self):
"""Builds the Q-value network computations needed for acting and training.
These are:
self.online_convnet: For computing the current state's Q-values.
self.target_convnet: For computing the next state's target Q-values.
self._net_outputs: The actual Q-values.
self._q_argmax: The action maximizing the current state's Q-values.
self._replay_net_outputs: The replayed states' Q-values.
self._replay_next_target_net_outputs: The replayed next states' target
Q-values (see Mnih et al., 2015 for details).
"""
# Calling online_convnet will generate a new graph as defined in
# self._get_network_template using whatever input is passed, but will always
# share the same weights.
self.online_convnet = tf.make_template('Online', self._network_template)
self.target_convnet = tf.make_template('Target', self._network_template)
self._net_outputs = self.online_convnet(self.state_ph)
# TODO(bellemare): Ties should be broken. They are unlikely to happen when
# using a deep network, but may affect performance with a linear
# approximation scheme.
self._q_argmax = tf.argmax(self._net_outputs.q_values, axis=1)[0]
self._replay_net_outputs = self.online_convnet(self._replay.states)
self._replay_next_target_net_outputs = self.target_convnet(
self._replay.next_states)
def __init__(self, trainable=False,
state_preprocess_net=lambda states: states,
action_embed_net=lambda actions, *args, **kwargs: actions,
ndims=None):
self.trainable = trainable
self._scope = tf.get_variable_scope().name
self._ndims = ndims
self._state_preprocess_net = tf.make_template(
self.STATE_PREPROCESS_NET_SCOPE, state_preprocess_net,
create_scope_now_=True)
self._action_embed_net = tf.make_template(
self.ACTION_EMBED_NET_SCOPE, action_embed_net,
create_scope_now_=True)
def __init__(self, corpus, **opts):
self.corpus = corpus
self.opts = opts
self.global_step = get_or_create_global_step()
self.increment_global_step_op = tf.assign(self.global_step, self.global_step + 1, name="increment_global_step")
self.corpus_size = get_corpus_size(self.corpus["train"])
self.corpus_size_valid = get_corpus_size(self.corpus["valid"])
self.word2idx, self.idx2word = build_vocab(self.corpus["train"])
self.vocab_size = len(self.word2idx)
self.generator_template = tf.make_template(GENERATOR_PREFIX, generator)
self.discriminator_template = tf.make_template(DISCRIMINATOR_PREFIX, discriminator)
self.enqueue_data, _, source, target, sequence_length = \
prepare_data(self.corpus["train"], self.word2idx, num_threads=7, **self.opts)
# TODO: option to either do pretrain or just generate?
self.g_tensors_pretrain = self.generator_template(
source, target, sequence_length, self.vocab_size, **self.opts)
self.enqueue_data_valid, self.input_ph, source_valid, target_valid, sequence_length_valid = \
prepare_data(self.corpus["valid"], self.word2idx, num_threads=1, **self.opts)
self.g_tensors_pretrain_valid = self.generator_template(
source_valid, target_valid, sequence_length_valid, self.vocab_size, **self.opts)
self.decoder_fn = prepare_custom_decoder(
sequence_length, self.g_tensors_pretrain.embedding_matrix, self.g_tensors_pretrain.output_projections)
self.g_tensors_fake = self.generator_template(
source, target, sequence_length, self.vocab_size, decoder_fn=self.decoder_fn, **self.opts)
self.g_tensors_fake_valid = self.generator_template(
source_valid, target_valid, sequence_length_valid, self.vocab_size, decoder_fn=self.decoder_fn, **self.opts)
# TODO: using the rnn outputs from pretraining as "real" instead of target embeddings (aka professor forcing)
self.d_tensors_real = self.discriminator_template(
self.g_tensors_pretrain.rnn_outputs, sequence_length, is_real=True, **self.opts)
# TODO: check to see if sequence_length is correct
self.d_tensors_fake = self.discriminator_template(
self.g_tensors_fake.rnn_outputs, None, is_real=False, **self.opts)
self.g_tvars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope=GENERATOR_PREFIX)
self.d_tvars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope=DISCRIMINATOR_PREFIX)
def _initialize_policy(self):
"""Initialize the policy.
Run the policy network on dummy data to initialize its parameters for later
reuse and to analyze the policy distribution. Initializes the attributes
`self._network` and `self._policy_type`.
Raises:
ValueError: Invalid policy distribution.
Returns:
Parameters of the policy distribution and policy state.
"""
with tf.device('/gpu:0' if self._use_gpu else '/cpu:0'):
network = functools.partial(
self._config.network, self._config, self._batch_env.action_space)
self._network = tf.make_template('network', network)
output = self._network(
tf.zeros_like(self._batch_env.observ)[:, None],
tf.ones(len(self._batch_env)))
if output.policy.event_shape != self._batch_env.action.shape[1:]:
message = 'Policy event shape {} does not match action shape {}.'
message = message.format(
output.policy.event_shape, self._batch_env.action.shape[1:])
raise ValueError(message)
self._policy_type = type(output.policy)
is_tensor = lambda x: isinstance(x, tf.Tensor)
policy_params = tools.nested.filter(is_tensor, output.policy.parameters)
set_batch_dim = lambda x: utility.set_dimension(x, 0, len(self._batch_env))
tools.nested.map(set_batch_dim, policy_params)
if output.state is not None:
tools.nested.map(set_batch_dim, output.state)
return policy_params, output.state
def __init__(self, *args, **kwargs):
self.func = func
self.args = args
self.kwargs = kwargs
self.name = self.kwargs.get("name", self.func.__name__)
self._template = tf.make_template(self.name, self.func, create_scope_now_=True)
self._unique_name = self._template.variable_scope.name.split("/")[-1]
self._summary_added = False
def test_variable_reuse_with_template(self):
tmpl1 = tf.make_template("test", tf.contrib.layers.legacy_fully_connected, num_output_units=8)
output1 = tmpl1(self.input)
output2 = tmpl1(self.input)
with tf.Session() as sess:
tf.initialize_all_variables().run()
out_value1, out_value2 = sess.run([output1, output2])
self.assertAllClose(out_value1, out_value2)
def build_model(self):
sc = predictron_arg_scope()
with tf.variable_scope('state'):
with slim.arg_scope(sc):
state = slim.conv2d(self.inputs, 32, [3, 3], scope='conv1')
state = layers.batch_norm(state, activation_fn=tf.nn.relu, scope='conv1/preact')
state = slim.conv2d(state, 32, [3, 3], scope='conv2')
state = layers.batch_norm(state, activation_fn=tf.nn.relu, scope='conv2/preact')
iter_template = tf.make_template('iter', self.iter_func, unique_name_='iter')
rewards_arr = []
gammas_arr = []
lambdas_arr = []
values_arr = []
for k in range(self.max_depth):
state, reward, gamma, lambda_, value = iter_template(state)
rewards_arr.append(reward)
gammas_arr.append(gamma)
lambdas_arr.append(lambda_)
values_arr.append(value)
_, _, _, _, value = iter_template(state)
# K + 1 elements
values_arr.append(value)
bs = tf.shape(self.inputs)[0]
# [batch_size, K * maze_size]
self.rewards = tf.pack(rewards_arr, axis=1)
# [batch_size, K, maze_size]
self.rewards = tf.reshape(self.rewards, [bs, self.max_depth, self.maze_size])
# [batch_size, K + 1, maze_size]
self.rewards = tf.concat_v2(values=[tf.zeros(shape=[bs, 1, self.maze_size], dtype=tf.float32), self.rewards],
axis=1, name='rewards')
# [batch_size, K * maze_size]
self.gammas = tf.pack(gammas_arr, axis=1)
# [batch_size, K, maze_size]
self.gammas = tf.reshape(self.gammas, [bs, self.max_depth, self.maze_size])
# [batch_size, K + 1, maze_size]
self.gammas = tf.concat_v2(values=[tf.ones(shape=[bs, 1, self.maze_size], dtype=tf.float32), self.gammas],
axis=1, name='gammas')
# [batch_size, K * maze_size]
self.lambdas = tf.pack(lambdas_arr, axis=1)
# [batch_size, K, maze_size]
self.lambdas = tf.reshape(self.lambdas, [-1, self.max_depth, self.maze_size])
# [batch_size, (K + 1) * maze_size]
self.values = tf.pack(values_arr, axis=1)
# [batch_size, K + 1, maze_size]
self.values = tf.reshape(self.values, [-1, (self.max_depth + 1), self.maze_size])
self.build_preturns()
self.build_lambda_preturns()
def test_variable_reuse_with_template(self):
tmpl1 = tf.make_template("test", tf.learn.fully_connected, num_output_nodes=8)
output1 = tmpl1(self.input)
output2 = tmpl1(self.input)
with tf.Session() as sess:
tf.initialize_all_variables().run()
out_value1, out_value2 = sess.run([output1, output2])
self.assertAllClose(out_value1, out_value2)
assert_summary_scope(r"test(_\d)?/fully_connected")
def test_variable_reuse_with_template(self):
tmpl1 = tf.make_template("test", tf.learn.convolution2d, kernel_size=(3, 3), num_output_channels=8)
output1 = tmpl1(self.input)
output2 = tmpl1(self.input)
with tf.Session() as sess:
tf.initialize_all_variables().run()
out_value1, out_value2 = sess.run([output1, output2])
self.assertAllClose(out_value1, out_value2)
assert_summary_scope(r"test(_\d)?/convolution2d")
def testBijectorConditionKwargs(self):
batch_size = 3
x_ = np.linspace(-1.0, 1.0, (batch_size * 4 * 2)).astype(
np.float32).reshape((batch_size, 4 * 2))
conditions = {
"a": tf.random_normal((batch_size, 4), dtype=tf.float32),
"b": tf.random_normal((batch_size, 2), dtype=tf.float32),
}
def _condition_shift_and_log_scale_fn(x0, output_units, a, b):
x = tf.concat((x0, a, b), axis=-1)
out = tf.layers.dense(
inputs=x,
units=2 * output_units)
shift, log_scale = tf.split(out, 2, axis=-1)
return shift, log_scale
condition_shift_and_log_scale_fn = tf.make_template(
"real_nvp_condition_template", _condition_shift_and_log_scale_fn)
nvp = tfb.RealNVP(
num_masked=4,
validate_args=True,
is_constant_jacobian=False,
shift_and_log_scale_fn=condition_shift_and_log_scale_fn)
x = tf.constant(x_)
forward_x = nvp.forward(x, **conditions)
# Use identity to invalidate cache.
inverse_y = nvp.inverse(tf.identity(forward_x), **conditions)
forward_inverse_y = nvp.forward(inverse_y, **conditions)
fldj = nvp.forward_log_det_jacobian(x, event_ndims=1, **conditions)
# Use identity to invalidate cache.
ildj = nvp.inverse_log_det_jacobian(
tf.identity(forward_x), event_ndims=1, **conditions)
self.evaluate(tf.global_variables_initializer())
[
forward_x_,
inverse_y_,
forward_inverse_y_,
ildj_,
fldj_,
] = self.evaluate([
forward_x,
inverse_y,
forward_inverse_y,
ildj,
fldj,
])
self.assertEqual("real_nvp", nvp.name)
self.assertAllClose(forward_x_, forward_inverse_y_, rtol=1e-6, atol=0.)
self.assertAllClose(x_, inverse_y_, rtol=1e-6, atol=0.)
self.assertAllClose(ildj_, -fldj_, rtol=1e-6, atol=0.)
def initialize_graph(self, input_statistics):
"""Save templates for components, which can then be used repeatedly.
This method is called every time a new graph is created. It's safe to start
adding ops to the current default graph here, but the graph should be
constructed from scratch.
Args:
input_statistics: A math_utils.InputStatistics object.
"""
super(_LSTMModel, self).initialize_graph(input_statistics=input_statistics)
self._lstm_cell = tf.nn.rnn_cell.LSTMCell(num_units=self._num_units)
# Create templates so we don't have to worry about variable reuse.
self._lstm_cell_run = tf.make_template(
name_="lstm_cell",
func_=self._lstm_cell,
create_scope_now_=True)
# Transforms LSTM output into mean predictions.
self._predict_from_lstm_output = tf.make_template(
name_="predict_from_lstm_output",
func_=lambda inputs: tf.layers.dense(inputs=inputs, units=self.num_features),
create_scope_now_=True)
def test_variable_reuse_with_template(self):
tmpl1 = tf.make_template(
"test", tf.contrib.layers.legacy_convolution2d, kernel_size=(3, 3), num_output_channels=8
)
output1 = tmpl1(self.input)
output2 = tmpl1(self.input)
with tf.Session() as sess:
tf.initialize_all_variables().run()
out_value1, out_value2 = sess.run([output1, output2])
self.assertAllClose(out_value1, out_value2)
def __init__(self, mode=None, batch_size=hp_default.batch_size, queue=True):
self.mode = mode
self.batch_size = batch_size
self.queue = queue
self.is_training = self.get_is_training(mode)
# Input
self.x_mfcc, self.y_ppgs, self.y_spec, self.y_mel, self.num_batch = self.get_input(mode, batch_size, queue)
# Networks
self.net_template = tf.make_template('net', self._net2)
self.ppgs, self.pred_ppg, self.logits_ppg, self.pred_spec, self.pred_mel = self.net_template()
def __init__(self,
f,
g,
num_layers=1,
f_side_input=None,
g_side_input=None,
use_efficient_backprop=True):
if isinstance(f, list):
assert len(f) == num_layers
else:
f = [f] * num_layers
if isinstance(g, list):
assert len(g) == num_layers
else:
g = [g] * num_layers
scope_prefix = "revblock/revlayer_%d/"
f_scope = scope_prefix + "f"
g_scope = scope_prefix + "g"
f = [
tf.make_template(f_scope % i, fn, create_scope_now_=True)
for i, fn in enumerate(f)
]
g = [
tf.make_template(g_scope % i, fn, create_scope_now_=True)
for i, fn in enumerate(g)
]
self.f = f
self.g = g
self.num_layers = num_layers
self.f_side_input = f_side_input or []
self.g_side_input = g_side_input or []
self._use_efficient_backprop = use_efficient_backprop
def __init__(self, name):
"""
Initialize the module. Each subclass must call this constructor with a name.
Args:
name: Name of this module. Used for `tf.make_template`.
"""
self.name = name
self._template = tf.make_template(name, self._build, create_scope_now_=True)
# Docstrings for the class should be the docstring for the _build method
self.__doc__ = self._build.__doc__
# pylint: disable=E1101
self.__call__.__func__.__doc__ = self._build.__doc__
def test_all_ckpt(modelPath, fileOrDir,flags):
tf.reset_default_graph()
tf.logging.warning(modelPath)
tem = [f for f in os.listdir(modelPath) if 'data' in f]
ckptFiles = sorted([r.split('.data')[0] for r in tem])
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
with tf.Session(config=config) as sess:
input_tensor = tf.placeholder(tf.float32, shape=(1, None, None, 1))
shared_model = tf.make_template('shared_model', model)
output_tensor, weights = shared_model(input_tensor)
output_tensor = tf.clip_by_value(output_tensor, 0., 1.)
output_tensor = output_tensor * 255
saver = tf.train.Saver()
sess.run(tf.global_variables_initializer())
original_ycbcr, gt_y, fileName_list = prepare_test_data(fileOrDir)
for ckpt in ckptFiles:
epoch = int(ckpt.split('_')[-1].split('.')[0])
if flags==0:
if epoch != 555:
continue
elif flags==1:
if epoch!= 555:
continue
else:
if epoch != 555:
continue
tf.logging.warning("epoch:%d\t"%epoch)
saver.restore(sess,os.path.join(modelPath,ckpt))
total_imgs = len(fileName_list)
for i in range(total_imgs):
imgY = original_ycbcr[i][0]
out = sess.run(output_tensor, feed_dict={input_tensor: imgY})
out = np.reshape(out, (out.shape[1], out.shape[2]))
out = np.around(out)
out = out.astype('int')
out = out.tolist()
return out
def define_train(hparams, environment_spec, event_dir):
"""Define the training setup."""
if isinstance(environment_spec, str):
env_lambda = lambda: gym.make(environment_spec)
else:
env_lambda = environment_spec
policy_lambda = hparams.network
env = env_lambda()
action_space = env.action_space
batch_env = utils.define_batch_env(env_lambda, hparams.num_agents)
policy_factory = tf.make_template(
"network",
functools.partial(policy_lambda, action_space, hparams))
with tf.variable_scope("train"):
memory, collect_summary = collect.define_collect(
policy_factory, batch_env, hparams, eval_phase=False)
ppo_summary = ppo.define_ppo_epoch(memory, policy_factory, hparams)
summary = tf.summary.merge([collect_summary, ppo_summary])
with tf.variable_scope("eval"):
eval_env_lambda = env_lambda
if event_dir and hparams.video_during_eval:
# Some environments reset environments automatically, when reached done
# state. For them we shall record only every second episode.
d = 2 if env_lambda().metadata.get("semantics.autoreset") else 1
eval_env_lambda = lambda: gym.wrappers.Monitor( # pylint: disable=g-long-lambda
env_lambda(), event_dir, video_callable=lambda i: i % d == 0)
wrapped_eval_env_lambda = lambda: utils.EvalVideoWrapper(eval_env_lambda())
_, eval_summary = collect.define_collect(
policy_factory,
utils.define_batch_env(wrapped_eval_env_lambda, hparams.num_eval_agents,
xvfb=hparams.video_during_eval),
hparams, eval_phase=True)
return summary, eval_summary
def testMakeLogJointFnTemplate(self):
"""Test `make_log_joint_fn` on program returned by tf.make_template."""
def variational():
loc = tf.get_variable("loc", [])
qz = ed.Normal(loc=loc, scale=0.5, name="qz")
return qz
def true_log_joint(loc, qz):
log_prob = tf.reduce_sum(tfd.Normal(loc=loc, scale=0.5).log_prob(qz))
return log_prob
qz_value = 1.23
variational_template = tf.make_template("variational", variational)
log_joint = ed.make_log_joint_fn(variational_template)
expected_log_prob = log_joint(qz=qz_value)
loc = tf.trainable_variables("variational")[0]
actual_log_prob = true_log_joint(loc, qz_value)
with self.test_session() as sess:
sess.run(tf.initialize_all_variables())
actual_log_prob_, expected_log_prob_ = sess.run(
[actual_log_prob, expected_log_prob])
self.assertEqual(actual_log_prob_, expected_log_prob_)
def __init__(self, hps):
self.hps = hps
self.pixelcnn = tf.make_template('PixelCNN', pixelcnn_spec)
y_1 = ld.conv_layer(y_0, 3, 1, 8, "encode_3")
# conv5
conv5 = ld.transpose_conv_layer(y_1, 1, 1, 8, "decode_5")
# conv6
conv6 = ld.transpose_conv_layer(conv5, 3, 2, 8, "decode_6")
# conv7
conv7 = ld.transpose_conv_layer(conv6, 3, 1, 8, "decode_7")
# x_1
x_1 = ld.transpose_conv_layer(conv7, 3, 2, 5, "decode_8", True) # set activation to linear
return x_1, hidden
# make a template for reuse
network_template = tf.make_template('network', network)
def train():
"""Train ring_net for a number of steps."""
with tf.Graph().as_default():
# make inputs
x = tf.placeholder(tf.float32, [None, FLAGS.seq_length, 100, 100, 5])
# possible dropout inside
keep_prob = tf.placeholder("float")
x_dropout = tf.nn.dropout(x, keep_prob)
# create network
x_unwrap = []
def __init__(self, _sentinel=None, custom_getter=None, name=None): # pylint: disable=invalid-name
"""Performs the initialisation necessary for all AbstractModule instances.
Every subclass of AbstractModule must begin their constructor with a call to
this constructor, i.e. `super(MySubModule, self).__init__(name=name)`.
If you instantiate sub-modules in __init__ you must create them within the
`_enter_variable_scope` context manager to ensure they are in the module's
variable scope. Alternatively, instantiate sub-modules in `_build`.
Args:
_sentinel: Variable that only carries a non-None value if `__init__` was
called without named parameters. If this is the case, a deprecation
warning is issued in form of a `ValueError`.
custom_getter: Callable or dictionary of callables to use as
custom getters inside the module. If a dictionary, the keys
correspond to regexes to match variable names. See the `tf.get_variable`
documentation for information about the custom_getter API.
name: Name of this module. Used to construct the Templated build function.
If `None` the module's class name is used (converted to snake case).
Raises:
TypeError: If `name` is not a string.
TypeError: If a given `custom_getter` is not callable.
ValueError: If `__init__` was called without named arguments.
"""
if _sentinel is not None:
raise ValueError("Calling AbstractModule.__init__ without named "
"arguments is deprecated.")
if name is None:
name = util.to_snake_case(self.__class__.__name__)
elif not isinstance(name, six.string_types):
raise TypeError("Name must be a string.")
self._connected_subgraphs = []
# If the given custom getter is a dictionary with a per-variable custom
# getter, wrap it into a single custom getter.
if isinstance(custom_getter, collections.Mapping):
self._custom_getter = util._custom_getter_router( # pylint: disable=protected-access
custom_getter_map=custom_getter,
name_fn=lambda name: name[len(self.scope_name) + 1:])
else:
if not (custom_getter is None or callable(custom_getter)):
raise TypeError("Given custom_getter is not callable.")
self._custom_getter = custom_getter
self._template = tf.make_template(name,
self._build_wrapper,
create_scope_now_=True,
custom_getter_=self._custom_getter)
self._original_name = name
self._unique_name = self._template.variable_scope.name.split("/")[-1]
# Update __call__ and the object docstrings to enable better introspection
self.__doc__ = self._build.__doc__
self.__call__.__func__.__doc__ = self._build.__doc__
# Keep track of which graph this module has been connected to. Sonnet
# modules cannot be connected to multiple graphs, as transparent variable
# sharing is impossible in that case.
self._graph = None
def __init__(self):
def template_fn(x):
net = slim.fully_connected(x, 60)
return slim.fully_connected(net, 10, activation_fn=None)
self.template = tf.make_template('dummy_model', template_fn)
def define_model(data, trainer, config):
tf.logging.info('Build TensorFlow compute graph.')
dependencies = []
cleanups = []
step = trainer.step
global_step = trainer.global_step
phase = trainer.phase
# Instantiate network blocks.
cell = config.cell()
kwargs = dict(create_scope_now_=True)
encoder = tf.make_template('encoder', config.encoder, **kwargs)
embedded = encoder(data)
with tf.control_dependencies(dependencies):
embedded = tf.identity(embedded)
data['embedding'] = embedded
heads = tools.AttrDict(_unlocked=True)
dummy_features = cell.features_from_state(cell.zero_state(1, tf.float32))
for key, head in config.heads.items():
name = 'head_{}'.format(key)
kwargs = dict(create_scope_now_=True)
if key in data:
kwargs['data_shape'] = data[key].shape[2:].as_list()
elif key == 'action_target':
kwargs['data_shape'] = data['action'].shape[2:].as_list()
heads[key] = tf.make_template(name, head, **kwargs)
heads[key](dummy_features) # Initialize weights.
# Apply and optimize model.
graph = tools.AttrDict(locals())
prior, posterior = tools.unroll.closed_loop(cell, embedded, data['action'],
config.debug)
objectives = utility.compute_objectives(posterior, prior, data, graph,
config)
summaries, grad_norms = utility.apply_optimizers(objectives, trainer,
config)
# Active data collection.
with tf.variable_scope('collection'):
with tf.control_dependencies(summaries): # Make sure to train first.
for name, params in config.train_collects.items():
schedule = tools.schedule.binary(step, config.batch_shape[0],
params.steps_after,
params.steps_every,
params.steps_until)
summary, _ = tf.cond(
tf.logical_and(tf.equal(trainer.phase, 'train'), schedule),
functools.partial(utility.simulate_episodes,
config,
params,
graph,
cleanups,
expensive_summaries=False,
gif_summary=False,
name=name),
lambda: (tf.constant(''), tf.constant(0.0)),
name='should_collect_' + name)
summaries.append(summary)
# Compute summaries.
graph = tools.AttrDict(locals())
summary, score = tf.cond(
trainer.log,
lambda: define_summaries.define_summaries(graph, config, cleanups),
lambda: (tf.constant(''), tf.zeros((0, ), tf.float32)),
name='summaries')
summaries = tf.summary.merge([summaries, summary])
dependencies.append(
utility.print_metrics({ob.name: ob.value
for ob in objectives}, step,
config.print_metrics_every, 'objectives'))
dependencies.append(
utility.print_metrics(grad_norms, step, config.print_metrics_every,
'grad_norms'))
with tf.control_dependencies(dependencies):
score = tf.identity(score)
return score, summaries, cleanups
def real_nvp_default_template(hidden_layers,
shift_only=False,
activation=tf.nn.relu,
name=None,
*args, # pylint: disable=keyword-arg-before-vararg
**kwargs):
"""Build a scale-and-shift function using a multi-layer neural network.
This will be wrapped in a make_template to ensure the variables are only
created once. It takes the `d`-dimensional input x[0:d] and returns the `D-d`
dimensional outputs `loc` ("mu") and `log_scale` ("alpha").
The default template does not support conditioning and will raise an
exception if `condition_kwargs` are passed to it. To use conditioning in
real nvp bijector, implement a conditioned shift/scale template that
handles the `condition_kwargs`.
Arguments:
hidden_layers: Python `list`-like of non-negative integer, scalars
indicating the number of units in each hidden layer. Default: `[512, 512].
shift_only: Python `bool` indicating if only the `shift` term shall be
computed (i.e. NICE bijector). Default: `False`.
activation: Activation function (callable). Explicitly setting to `None`
implies a linear activation.
name: A name for ops managed by this function. Default:
"real_nvp_default_template".
*args: `tf.layers.dense` arguments.
**kwargs: `tf.layers.dense` keyword arguments.
Returns:
shift: `Float`-like `Tensor` of shift terms ("mu" in
[Papamakarios et al. (2016)][1]).
log_scale: `Float`-like `Tensor` of log(scale) terms ("alpha" in
[Papamakarios et al. (2016)][1]).
Raises:
NotImplementedError: if rightmost dimension of `inputs` is unknown prior to
graph execution, or if `condition_kwargs` is not empty.
#### References
[1]: George Papamakarios, Theo Pavlakou, and Iain Murray. Masked
Autoregressive Flow for Density Estimation. In _Neural Information
Processing Systems_, 2017. https://arxiv.org/abs/1705.07057
"""
with tf.name_scope(name, "real_nvp_default_template"):
def _fn(x, output_units, **condition_kwargs):
"""Fully connected MLP parameterized via `real_nvp_template`."""
if condition_kwargs:
raise NotImplementedError(
"Conditioning not implemented in the default template.")
for units in hidden_layers:
x = tf.layers.dense(
inputs=x,
units=units,
activation=activation,
*args, # pylint: disable=keyword-arg-before-vararg
**kwargs)
x = tf.layers.dense(
inputs=x,
units=(1 if shift_only else 2) * output_units,
activation=None,
*args, # pylint: disable=keyword-arg-before-vararg
**kwargs)
if shift_only:
return x, None
shift, log_scale = tf.split(x, 2, axis=-1)
return shift, log_scale
return tf.make_template("real_nvp_default_template", _fn)
def __init__(self, _sentinel=None, custom_getter=None,
name=None): # pylint: disable=invalid-name
"""Performs the initialisation necessary for all AbstractModule instances.
Every subclass of AbstractModule must begin their constructor with a call to
this constructor, i.e.
`super(MySubModule, self).__init__(custom_getter=custom_getter, name=name)`.
If you instantiate sub-modules in __init__ you must create them within the
`_enter_variable_scope` context manager to ensure they are in the module's
variable scope. Alternatively, instantiate sub-modules in `_build`.
Args:
_sentinel: Variable that only carries a non-None value if `__init__` was
called without named parameters. If this is the case, a deprecation
warning is issued in form of a `ValueError`.
custom_getter: Callable or dictionary of callables to use as
custom getters inside the module. If a dictionary, the keys
correspond to regexes to match variable names. See the `tf.get_variable`
documentation for information about the custom_getter API.
name: Name of this module. Used to construct the Templated build function.
If `None` the module's class name is used (converted to snake case).
Raises:
TypeError: If `name` is not a string.
TypeError: If a given `custom_getter` is not callable.
ValueError: If `__init__` was called without named arguments.
"""
if _sentinel is not None:
raise ValueError("Calling AbstractModule.__init__ without named "
"arguments is not supported.")
if name is None:
name = util.to_snake_case(self.__class__.__name__)
elif not isinstance(name, six.string_types):
raise TypeError("Name must be a string.")
self._connected_subgraphs = []
# If the given custom getter is a dictionary with a per-variable custom
# getter, wrap it into a single custom getter.
if isinstance(custom_getter, collections.Mapping):
self._custom_getter = util.custom_getter_router(
custom_getter_map=custom_getter,
name_fn=lambda name: name[len(self.scope_name) + 1:])
else:
if not (custom_getter is None or callable(custom_getter)):
raise TypeError("Given custom_getter is not callable.")
self._custom_getter = custom_getter
self._custom_getter = _maybe_wrap_custom_getter(
_variable_tracking_custom_getter, self._custom_getter)
self._template = tf.make_template(name,
self._build_wrapper,
create_scope_now_=True,
custom_getter_=self._custom_getter)
self._original_name = name
self._unique_name = self._template.variable_scope.name.split("/")[-1]
# Update __call__ and the object docstrings to enable better introspection.
self.__doc__ = self._build.__doc__
self.__call__.__func__.__doc__ = self._build.__doc__
# Keep track of which graph this module has been connected to. Sonnet
# modules cannot be connected to multiple graphs, as transparent variable
# sharing is impossible in that case.
self._graph = None
# Container for all variables created in this module and its sub-modules.
self._all_variables = set([])
def build_iteration(self,
iteration_number,
subnetwork_builders,
features,
mode,
labels=None,
previous_ensemble_summary=None,
previous_ensemble_spec=None,
rebuilding=False):
"""Builds and returns AdaNet iteration t.
This method uses the generated the candidate subnetworks given the ensemble
at iteration t-1 and creates graph operations to train them. The returned
`_Iteration` tracks the training of all candidates to know when the
iteration is over, and tracks the best candidate's predictions and loss, as
defined by lowest complexity-regularized loss on the train set.
Args:
iteration_number: Integer iteration number.
subnetwork_builders: A list of `Builders` for adding ` Subnetworks` to the
graph. Each subnetwork is then wrapped in a `_Candidate` to train.
features: Dictionary of `Tensor` objects keyed by feature name.
mode: Defines whether this is training, evaluation or prediction. See
`ModeKeys`.
labels: `Tensor` of labels. Can be `None`.
previous_ensemble_summary: The `_ScopedSummary` for the previous ensemble.
previous_ensemble_spec: Optional `_EnsembleSpec` for iteration t-1.
rebuilding: Boolean whether the iteration is being rebuilt only to restore
the previous best subnetworks and ensembles.
Returns:
An _Iteration instance.
Raises:
ValueError: If subnetwork_builders is empty.
ValueError: If two `Builder` instances share the same name.
"""
tf.logging.info("%s iteration %s",
"Rebuilding" if rebuilding else "Building",
iteration_number)
if not subnetwork_builders:
raise ValueError("Each iteration must have at least one Builder.")
# TODO: Consider moving ensemble mode logic to ensemble.py.
ensemble_mode = mode
if rebuilding:
# Create the frozen ensemble in EVAL mode by default. This way their
# outputs aren't affected by dropout etc.
ensemble_mode = tf.estimator.ModeKeys.EVAL
if mode == tf.estimator.ModeKeys.PREDICT:
ensemble_mode = mode
# Only replicate in training mode when the user requests it.
if self._replicate_ensemble_in_training and (
mode == tf.estimator.ModeKeys.TRAIN):
ensemble_mode = mode
training = mode == tf.estimator.ModeKeys.TRAIN
skip_summaries = mode == tf.estimator.ModeKeys.PREDICT
with tf.variable_scope("iteration_{}".format(iteration_number)):
# Iteration step to use instead of global step.
iteration_step = tf.get_variable(
"step",
shape=[],
initializer=tf.zeros_initializer(),
trainable=False,
dtype=tf.int64)
# Convert to tensor so that users cannot mutate it.
iteration_step_tensor = tf.convert_to_tensor(iteration_step)
seen_builder_names = {}
candidates = []
summaries = []
subnetwork_reports = {}
# TODO: Consolidate building subnetwork into
# candidate_builder.
if previous_ensemble_spec:
# Include previous best subnetwork as a candidate so that its
# predictions are returned until a new candidate outperforms.
seen_builder_names = {previous_ensemble_spec.name: True}
previous_best_candidate = self._candidate_builder.build_candidate(
ensemble_spec=previous_ensemble_spec,
training=training,
iteration_step=iteration_step_tensor,
summary=previous_ensemble_summary,
is_previous_best=True)
candidates.append(previous_best_candidate)
summaries.append(previous_ensemble_summary)
# Generate subnetwork reports.
if mode == tf.estimator.ModeKeys.EVAL:
subnetwork_report = subnetwork.Report(
hparams={},
attributes={},
metrics=(previous_ensemble_spec.eval_metric_ops.copy()
if previous_ensemble_spec.eval_metric_ops
is not None else {}),
)
subnetwork_report.metrics["adanet_loss"] = tf.metrics.mean(
previous_ensemble_spec.adanet_loss)
subnetwork_reports["previous_ensemble"] = subnetwork_report
for subnetwork_builder in subnetwork_builders:
if subnetwork_builder.name in seen_builder_names:
raise ValueError(
"Two ensembles have the same name '{}'".format(
subnetwork_builder.name))
seen_builder_names[subnetwork_builder.name] = True
ensemble_name = "t{}_{}".format(iteration_number,
subnetwork_builder.name)
summary = _ScopedSummary(ensemble_name,
skip_summary=skip_summaries
or rebuilding)
summaries.append(summary)
ensemble_spec = self._ensemble_builder.append_new_subnetwork(
ensemble_name=ensemble_name,
ensemble_spec=previous_ensemble_spec,
iteration_number=iteration_number,
subnetwork_builder=subnetwork_builder,
summary=summary,
features=features,
mode=ensemble_mode,
iteration_step=iteration_step_tensor,
labels=labels)
candidate = self._candidate_builder.build_candidate(
ensemble_spec=ensemble_spec,
training=training,
iteration_step=iteration_step_tensor,
summary=summary)
candidates.append(candidate)
# Generate subnetwork reports.
if mode != tf.estimator.ModeKeys.PREDICT:
subnetwork_report = subnetwork_builder.build_subnetwork_report(
)
if not subnetwork_report:
subnetwork_report = subnetwork.Report(hparams={},
attributes={},
metrics={})
if ensemble_spec.eval_metric_ops is not None:
for metric_name in sorted(
ensemble_spec.eval_metric_ops):
metric = ensemble_spec.eval_metric_ops[metric_name]
subnetwork_report.metrics[metric_name] = metric
subnetwork_report.metrics["adanet_loss"] = tf.metrics.mean(
ensemble_spec.adanet_loss)
subnetwork_reports[
subnetwork_builder.name] = subnetwork_report
# Dynamically select the outputs of best candidate.
best_candidate_index = self._best_candidate_index(candidates)
best_predictions = self._best_predictions(candidates,
best_candidate_index)
best_loss = self._best_loss(candidates, best_candidate_index, mode)
best_eval_metric_ops = self._best_eval_metric_ops(
candidates, best_candidate_index, mode)
best_export_outputs = self._best_export_outputs(
candidates, best_candidate_index, mode, best_predictions)
# Hooks on TPU cannot depend on any graph `Tensors`. Instead the value of
# `is_over` is stored in a `Variable` that can later be retrieved from
# inside a training hook.
is_over_var_fn = tf.make_template("is_over_var_fn", is_over_var)
estimator_spec = tf.estimator.EstimatorSpec(
mode=mode,
predictions=best_predictions,
loss=best_loss,
train_op=self._create_train_op(candidates, mode,
iteration_step, is_over_var_fn),
eval_metric_ops=best_eval_metric_ops,
export_outputs=best_export_outputs)
return _Iteration(number=iteration_number,
candidates=candidates,
estimator_spec=estimator_spec,
best_candidate_index=best_candidate_index,
summaries=summaries,
is_over_fn=is_over_var_fn,
subnetwork_reports=subnetwork_reports,
step=iteration_step_tensor)
def __init__(self, sess, args, seed, input_shape, n_train_tasks):
random.seed(seed)
np.random.seed(seed)
tf.set_random_seed(seed)
# get parsed args
self.lr = args.lr
self.batch_size = args.batch_size
self.sess = sess
self.summary = False
self.summary_dir = args.summary_dir
if (self.summary_dir):
self.summary = True
self.summary_interval = 200
summaries_list_train = []
summaries_list_val_test = []
summaries_list_finetune = []
summaries_list_test = []
else:
self.summary_dir = "no_summary"
self.val_task_finetune_interval = 500
self.val_task_finetune_epochs = 1000
self.early_stopping_val = 100
self.finetune_data_percentage = 0.8
self.n_classes = 1
self.input_height, self.input_width = 28, 28
self.n_train_tasks = n_train_tasks
self.flatten = tf.keras.layers.Flatten()
if (args.filters == ""):
self.filter_sizes = []
else:
self.filter_sizes = [int(i) for i in args.filters.split(' ')]
self.kernel_sizes = [int(i) for i in args.kernel_sizes.split(' ')]
# build model
self.shared_layers = []
if (len(self.filter_sizes) > 0):
self.shared_layers.append(
tf.keras.layers.Conv2D(filters=self.filter_sizes[0],
kernel_size=self.kernel_sizes[0],
input_shape=(None, self.input_height,
self.input_width),
strides=2,
padding='same',
activation='relu',
name='conv0_shared'))
self.shared_layers.append(
tf.keras.layers.BatchNormalization(name='bn0_shared'))
for i in range(1, len(self.filter_sizes)):
self.shared_layers.append(
tf.keras.layers.Conv2D(filters=self.filter_sizes[i],
kernel_size=self.kernel_sizes[i],
strides=2,
padding='same',
activation='relu',
name='conv' + str(i) + '_shared'))
self.shared_layers.append(
tf.keras.layers.BatchNormalization(name='bn' + str(i) +
'_shared'))
else:
for i in range(0, len(self.dense_sizes) - 2):
self.shared_layers.append(
tf.keras.layers.Dense(units=self.dense_sizes[i],
activation='relu',
name='dense_shared_' + str(i)))
if (args.dense_layers == ""):
self.dense_sizes = []
elif (' ' not in args.dense_layers):
self.dense_sizes = []
self.dense_sizes.append(int(args.dense_layers))
else:
self.dense_sizes = [int(i) for i in args.dense_layers.split(' ')]
self.task_layers = []
for i in range(self.n_train_tasks + 2):
for j in range(0, len(self.dense_sizes)):
self.task_layers.append(
tf.keras.layers.Dense(units=self.dense_sizes[j],
activation='relu',
name='dense' + str(j) + '_separate' +
str(i)))
self.task_layers.append(
tf.keras.layers.BatchNormalization(name='bn' + str(j) +
'_separate' + str(i)))
self.task_layers.append(
tf.keras.layers.Dense(units=self.n_classes,
name='output_layer' + str(i)))
# loss function
self.loss_fct = tf.nn.sigmoid_cross_entropy_with_logits
self.X_train = tf.placeholder(
tf.float32, (None, None, self.input_height, self.input_width),
name='X_train')
self.Y_train = tf.placeholder(tf.float32, (None, None, self.n_classes),
name='Y_train')
self.X_finetune = tf.placeholder(
tf.float32, (None, self.input_height, self.input_width),
name='X_finetune')
self.Y_finetune = tf.placeholder(tf.float32, (None, self.n_classes),
name='Y_finetune')
self.construct_forward_shared = tf.make_template(
'construct_forward_shared', self.feed_forward_shared)
self.train_losses = self.compute_losses(update_batchnorm=True)
shared_vars = tf.global_variables(scope="construct_forward_shared/")
bn_update_ops_shared_train = []
for layer in self.shared_layers:
if ('bn' in layer.name):
bn_update_ops_shared_train.append(layer.updates)
bn_update_ops_train_task = []
n_denses_bn = len(self.dense_sizes) * 2 + 1
for layer in self.task_layers[:-2 * n_denses_bn]:
if ('bn' in layer.name):
bn_update_ops_train_task.append(layer.updates)
self.train_update_ops = []
self.train_update_ops += bn_update_ops_shared_train
self.train_update_ops += bn_update_ops_train_task
for i in range(self.n_train_tasks):
task_vars = tf.global_variables(scope="task" + str(i) + "/")
self.train_update_ops.append(
tf.train.AdamOptimizer(self.lr).minimize(self.train_losses[i],
var_list=shared_vars +
task_vars))
self.mean_train_loss = tf.reduce_mean(self.train_losses)
if (self.summary):
summaries_list_train.append(
tf.summary.scalar('mean_train_loss', self.mean_train_loss))
self.merged_train = tf.summary.merge(summaries_list_train)
val_finetune_shared_out = self.construct_forward_shared(
self.X_finetune, True)
val_test_shared_out = self.construct_forward_shared(
self.X_finetune, False)
with tf.variable_scope('task8_val', reuse=tf.AUTO_REUSE):
self.val_finetune_output = self.feed_forward_task(
val_finetune_shared_out, True, -2)
self.val_finetune_loss = tf.reduce_mean(
self.loss_fct(labels=self.Y_finetune,
logits=self.val_finetune_output))
bn_update_ops_shared_val_finetune = []
for layer in self.shared_layers:
if ('bn' in layer.name):
bn_update_ops_shared_val_finetune.append(layer.updates[-2:])
bn_update_ops_finetune_val_task = []
for layer in self.task_layers[-2 * n_denses_bn:-n_denses_bn]:
if ('bn' in layer.name):
bn_update_ops_finetune_val_task.append(layer.updates)
self.val_finetune_update_op = []
self.val_finetune_update_op += bn_update_ops_shared_val_finetune
self.val_finetune_update_op += bn_update_ops_finetune_val_task
finetune_val_task_vars = tf.global_variables(scope="task8_val/")
self.val_finetune_update_op.append(
tf.train.AdamOptimizer(self.lr).minimize(
self.val_finetune_loss,
var_list=finetune_val_task_vars + shared_vars))
self.val_test_output = self.feed_forward_task(val_test_shared_out,
False, -2)
self.val_test_loss = tf.reduce_mean(
self.loss_fct(labels=self.Y_finetune, logits=self.val_test_output))
self.val_test_acc, self.val_test_precision, self.val_test_recall, self.val_test_specificity, self.val_test_f1_score, self.val_test_auc_pr = self.compute_metrics(
self.val_test_output, self.Y_finetune)
finetune_shared_out = self.construct_forward_shared(
self.X_finetune, True)
test_shared_out = self.construct_forward_shared(self.X_finetune, False)
with tf.variable_scope('task9_test', reuse=tf.AUTO_REUSE):
self.finetune_output = self.feed_forward_task(
finetune_shared_out, True, -1)
self.finetune_loss = tf.reduce_mean(
self.loss_fct(labels=self.Y_finetune, logits=self.finetune_output))
bn_update_ops_shared_test_finetune = []
for layer in self.shared_layers:
if ('bn' in layer.name):
bn_update_ops_shared_test_finetune.append(layer.updates[-2:])
bn_update_ops_finetune_test_task = []
for layer in self.task_layers[-n_denses_bn:]:
if ('bn' in layer.name):
bn_update_ops_finetune_test_task.append(layer.updates)
self.finetune_update_op = []
self.finetune_update_op += bn_update_ops_shared_test_finetune
self.finetune_update_op += bn_update_ops_finetune_test_task
finetune_task_vars = tf.global_variables(scope="task9_test/")
self.finetune_update_op.append(
tf.train.AdamOptimizer(self.lr).minimize(
self.finetune_loss, var_list=finetune_task_vars + shared_vars))
self.test_output = self.feed_forward_task(test_shared_out, False, -1)
self.test_loss = tf.reduce_mean(
self.loss_fct(labels=self.Y_finetune, logits=self.test_output))
self.test_acc, self.test_precision, self.test_recall, self.test_specificity, self.test_f1_score, self.test_auc_pr = self.compute_metrics(
self.test_output, self.Y_finetune)
if (self.summary):
summaries_list_finetune.append(
tf.summary.scalar('finetune_loss', self.finetune_loss))
summaries_list_test.append(
tf.summary.scalar('test_loss', self.test_loss))
summaries_list_test.append(
tf.summary.scalar('test_acc', self.test_acc))
summaries_list_test.append(
tf.summary.scalar('test_precision', self.test_precision))
summaries_list_test.append(
tf.summary.scalar('test_recall', self.test_recall))
summaries_list_test.append(
tf.summary.scalar('test_specificity', self.test_specificity))
summaries_list_test.append(
tf.summary.scalar('test_f1_score', self.test_f1_score))
summaries_list_test.append(
tf.summary.scalar('test_auc_pr', self.test_auc_pr))
self.merged_finetune = tf.summary.merge(summaries_list_finetune)
self.merged_test = tf.summary.merge(summaries_list_test)
summaries_list_val_test.append(
tf.summary.scalar('val_test_loss', self.val_test_loss))
summaries_list_val_test.append(
tf.summary.scalar('val_test_acc', self.val_test_acc))
summaries_list_val_test.append(
tf.summary.scalar('val_test_precision',
self.val_test_precision))
summaries_list_val_test.append(
tf.summary.scalar('val_test_recall', self.val_test_recall))
summaries_list_val_test.append(
tf.summary.scalar('val_test_specificity',
self.val_test_specificity))
summaries_list_val_test.append(
tf.summary.scalar('val_test_f1_score', self.val_test_f1_score))
summaries_list_val_test.append(
tf.summary.scalar('val_test_auc_pr', self.val_test_auc_pr))
self.merged_val_test = tf.summary.merge(summaries_list_val_test)
self.saver = tf.train.Saver(max_to_keep=250)
base_path = '/home/USER/Documents'
if (not (os.path.exists(base_path))):
base_path = '/home/ubuntu/Projects'
self.checkpoint_path = base_path + '/MAML/checkpoints_MTL/'
if (not (os.path.exists(self.checkpoint_path))):
os.mkdir(self.checkpoint_path)
if (not (os.path.exists(
os.path.join(self.checkpoint_path, self.summary_dir)))):
os.mkdir(os.path.join(self.checkpoint_path, self.summary_dir))
def fit_with_hmc(model,
observed_time_series,
num_results=100,
num_warmup_steps=50,
num_leapfrog_steps=15,
initial_state=None,
initial_step_size=None,
chain_batch_shape=(),
num_variational_steps=150,
variational_optimizer=None,
seed=None,
name=None):
"""Draw posterior samples using Hamiltonian Monte Carlo (HMC).
Markov chain Monte Carlo (MCMC) methods are considered the gold standard of
Bayesian inference; under suitable conditions and in the limit of infinitely
many draws they generate samples from the true posterior distribution. HMC [1]
uses gradients of the model's log-density function to propose samples,
allowing it to exploit posterior geometry. However, it is computationally more
expensive than variational inference and relatively sensitive to tuning.
This method attempts to provide a sensible default approach for fitting
StructuralTimeSeries models using HMC. It first runs variational inference as
a fast posterior approximation, and initializes the HMC sampler from the
variational posterior, using the posterior standard deviations to set
per-variable step sizes (equivalently, a diagonal mass matrix). During the
warmup phase, it adapts the step size to target an acceptance rate of 0.75,
which is thought to be in the desirable range for optimal mixing [2].
Args:
model: An instance of `StructuralTimeSeries` representing a
time-series model. This represents a joint distribution over
time-series and their parameters with batch shape `[b1, ..., bN]`.
observed_time_series: `float` `Tensor` of shape
`concat([sample_shape, model.batch_shape, [num_timesteps, 1]]) where
`sample_shape` corresponds to i.i.d. observations, and the trailing `[1]`
dimension may (optionally) be omitted if `num_timesteps > 1`.
num_results: Integer number of Markov chain draws.
Default value: `100`.
num_warmup_steps: Integer number of steps to take before starting to
collect results. The warmup steps are also used to adapt the step size
towards a target acceptance rate of 0.75.
Default value: `50`.
num_leapfrog_steps: Integer number of steps to run the leapfrog integrator
for. Total progress per HMC step is roughly proportional to
`step_size * num_leapfrog_steps`.
Default value: `15`.
initial_state: Optional Python `list` of `Tensor`s, one for each model
parameter, representing the initial state(s) of the Markov chain(s). These
should have shape `concat([chain_batch_shape, param.prior.batch_shape,
param.prior.event_shape])`. If `None`, the initial state is set
automatically using a sample from a variational posterior.
Default value: `None`.
initial_step_size: Python `list` of `Tensor`s, one for each model parameter,
representing the step size for the leapfrog integrator. Must
broadcast with the shape of `initial_state`. Larger step sizes lead to
faster progress, but too-large step sizes make rejection exponentially
more likely. If `None`, the step size is set automatically using the
standard deviation of a variational posterior.
Default value: `None`.
chain_batch_shape: Batch shape (Python `tuple`, `list`, or `int`) of chains
to run in parallel.
Default value: `[]` (i.e., a single chain).
num_variational_steps: Python `int` number of steps to run the variational
optimization to determine the initial state and step sizes.
Default value: `200`.
variational_optimizer: Optional `tf.train.Optimizer` instance to use in
the variational optimization. If `None`, defaults to
`tf.train.AdamOptimizer(0.1)`.
Default value: `None`.
seed: Python integer to seed the random number generator.
name: Python `str` name prefixed to ops created by this function.
Default value: `None` (i.e., 'fit_with_hmc').
Returns:
samples: Python `list` of `Tensors` representing posterior samples of model
parameters, with shapes `[concat([[num_results], chain_batch_shape,
param.prior.batch_shape, param.prior.event_shape]) for param in
model.parameters]`.
kernel_results: A (possibly nested) `tuple`, `namedtuple` or `list` of
`Tensor`s representing internal calculations made within the HMC sampler.
#### Examples
Assume we've built a structural time-series model:
```python
day_of_week = tfp.sts.Seasonal(
num_seasons=7,
observed_time_series=observed_time_series,
name='day_of_week')
local_linear_trend = tfp.sts.LocalLinearTrend(
observed_time_series=observed_time_series,
name='local_linear_trend')
model = tfp.sts.Sum(components=[day_of_week, local_linear_trend],
observed_time_series=observed_time_series)
```
To draw posterior samples using HMC under default settings:
```python
samples, kernel_results = tfp.sts.fit_with_hmc(model, observed_time_series)
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
samples_, kernel_results_ = sess.run((samples, kernel_results))
print("acceptance rate: {}".format(
np.mean(kernel_results_.inner_results.is_accepted, axis=0)))
print("posterior means: {}".format(
{param.name: np.mean(param_draws, axis=0)
for (param, param_draws) in zip(model.parameters, samples_)}))
```
We can also run multiple chains. This may help diagnose convergence issues
and allows us to exploit vectorization to draw samples more quickly, although
warmup still requires the same number of sequential steps.
```python
from matplotlib import pylab as plt
samples, kernel_results = tfp.sts.fit_with_hmc(
model, observed_time_series, chain_batch_shape=[10])
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
samples_, kernel_results_ = sess.run((samples, kernel_results))
print("acceptance rate: {}".format(
np.mean(kernel_results_.inner_results.is_accepted, axis=0)))
# Plot the sampled traces for each parameter. If the chains have mixed, their
# traces should all cover the same region of state space, frequently crossing
# over each other.
for (param, param_draws) in zip(model.parameters, samples_):
if param.prior.event_shape.ndims > 0:
print("Only plotting traces for scalar parameters, skipping {}".format(
param.name))
continue
plt.figure(figsize=[10, 4])
plt.title(param.name)
plt.plot(param_draws)
plt.ylabel(param.name)
plt.xlabel("HMC step")
# Combining the samples from multiple chains into a single dimension allows
# us to easily pass sampled parameters to downstream forecasting methods.
combined_samples_ = [np.reshape(param_draws,
[-1] + list(param_draws.shape[2:]))
for param_draws in samples_]
```
For greater flexibility, you may prefer to implement your own sampler using
the TensorFlow Probability primitives in `tfp.mcmc`. The following recipe
constructs a basic HMC sampler, using a `TransformedTransitionKernel` to
incorporate constraints on the parameter space.
```python
transformed_hmc_kernel = mcmc.TransformedTransitionKernel(
inner_kernel=mcmc.HamiltonianMonteCarlo(
target_log_prob_fn=model.joint_log_prob(observed_time_series),
step_size=step_size,
num_leapfrog_steps=num_leapfrog_steps,
step_size_update_fn=tfp.mcmc.make_simple_step_size_update_policy(
num_adaptation_steps=num_adaptation_steps),
state_gradients_are_stopped=True,
seed=seed),
bijector=[param.bijector for param in model.parameters])
# Initialize from a Uniform[-2, 2] distribution in unconstrained space.
initial_state = [tfp.sts.sample_uniform_initial_state(
param, return_constrained=True) for param in model.parameters]
samples, kernel_results = tfp.mcmc.sample_chain(
kernel=transformed_hmc_kernel,
num_results=num_results,
current_state=initial_state,
num_burnin_steps=num_warmup_steps)
```
#### References
[1]: Radford Neal. MCMC Using Hamiltonian Dynamics. _Handbook of Markov Chain
Monte Carlo_, 2011. https://arxiv.org/abs/1206.1901
[2] M.J. Betancourt, Simon Byrne, and Mark Girolami. Optimizing The
Integrator Step Size for Hamiltonian Monte Carlo.
https://arxiv.org/abs/1411.6669
"""
with tf.name_scope(name, 'fit_with_hmc',
values=[observed_time_series]) as name:
observed_time_series = tf.convert_to_tensor(observed_time_series,
name='observed_time_series')
seed = tfd.SeedStream(seed, salt='StructuralTimeSeries_fit_with_hmc')
# Initialize state and step sizes from a variational posterior if not
# specified.
if initial_step_size is None or initial_state is None:
# To avoid threading variational distributions through the training
# while loop, we build our own copy here. `make_template` ensures
# that our variational distributions share the optimized parameters.
def make_variational():
return build_factored_variational_loss(
model, observed_time_series,
init_batch_shape=chain_batch_shape, seed=seed())
make_variational = tf.make_template('make_variational', make_variational)
_, variational_distributions = make_variational()
minimize_op = _minimize_in_graph(
build_loss_fn=lambda: make_variational()[0], # return just the loss.
num_steps=num_variational_steps,
optimizer=variational_optimizer)
with tf.control_dependencies([minimize_op]):
if initial_state is None:
initial_state = [tf.stop_gradient(d.sample())
for d in variational_distributions.values()]
# Set step sizes using the unconstrained variational distribution.
if initial_step_size is None:
initial_step_size = [
transformed_q.distribution.stddev()
for transformed_q in variational_distributions.values()]
# Multiple chains manifest as an extra param batch dimension, so we need to
# add a corresponding batch dimension to `observed_time_series`.
observed_time_series = pad_batch_dimension_for_multiple_chains(
observed_time_series, model, chain_batch_shape=chain_batch_shape)
# When the initial step size depends on a variational optimization, we
# can't initialize step size variables before the optimization runs.
# Instead we initialize with a dummy value of the appropriate
# shape, then wrap the HMC chain with `control_dependencies` to ensure the
# variational step sizes are assigned before HMC actually runs.
step_size = [tf.get_variable(
initializer=tf.zeros_like(sample_uniform_initial_state(
param, init_sample_shape=chain_batch_shape,
return_constrained=False)),
name='{}_step_size'.format(param.name),
trainable=False,
use_resource=True)
for (param, ss) in zip(model.parameters, initial_step_size)]
step_size_init_op = tf.group(
[tf.assign(ss, initial_ss)
for (ss, initial_ss) in zip(step_size, initial_step_size)])
# Run HMC to sample from the posterior on parameters.
with tf.control_dependencies([step_size_init_op]):
samples, kernel_results = mcmc.sample_chain(
num_results=num_results,
current_state=initial_state,
num_burnin_steps=num_warmup_steps,
kernel=mcmc.TransformedTransitionKernel(
inner_kernel=mcmc.HamiltonianMonteCarlo(
target_log_prob_fn=model.joint_log_prob(observed_time_series),
step_size=step_size,
num_leapfrog_steps=num_leapfrog_steps,
step_size_update_fn=mcmc.make_simple_step_size_update_policy(
num_adaptation_steps=int(num_warmup_steps * 0.8),
decrement_multiplier=0.1,
increment_multiplier=0.1),
state_gradients_are_stopped=True,
seed=seed()),
bijector=[param.bijector for param in model.parameters]),
parallel_iterations=1 if seed is not None else 10)
return samples, kernel_results
def _build_all_bisimulation_parts(self):
"""Builds the bisimulation networks and ops."""
self.batch_size = tf.shape(self._replay.rewards)[0]
self._replay_target_outputs = self.target_convnet(self._replay.states)
self.bisim_horizon_ph = tf.placeholder(tf.float32, ())
self.online_bisimulation = tf.make_template('OnlineBisim',
bisimulation_network)
self.target_bisimulation = tf.make_template('TargetBisim',
bisimulation_network,
trainable=False)
# For evaluating the metric from an episode's first state.
self.source_state_ph = tf.placeholder(self.observation_dtype,
self.state_ph.shape,
name='source_state_ph')
self._initial_state_net = self.online_convnet(self.source_state_ph)
concat_states = tf.concat([
self._initial_state_net.representation,
self._net_outputs.representation
], 1)
self.state_distances = tf.squeeze(
self.online_bisimulation(concat_states))
self.state_value = tf.reduce_max(self._net_outputs.q_values, axis=1)[0]
if self.summary_writer is not None:
tf.summary.scalar('Eval/StateDistances', self.state_distances)
if self.evaluate_metric_only:
return
self.s1_online_distances = self.online_bisimulation(
self._concat_states(self._replay_net_outputs.representation))
self.s2_target_distances = self.target_bisimulation(
self._concat_states(
self._replay_next_target_net_outputs.representation))
# bisimulation_target = rew_diff + gamma * next_distance.
bisimulation_target = tf.stop_gradient(
self._build_bisimulation_target())
# We zero-out diagonal entries, since those are estimating the distance
# between a state and itself, which we know to be 0.
diagonal_mask = 1.0 - tf.diag(
tf.ones(self.batch_size, dtype=tf.float32))
diagonal_mask = tf.reshape(diagonal_mask, (self.batch_size**2, 1))
bisimulation_target *= diagonal_mask
bisimulation_estimate = self.s1_online_distances
bisimulation_loss = tf.losses.mean_squared_error(
bisimulation_target, bisimulation_estimate)
if self.summary_writer is not None:
average_distance = tf.reduce_mean(bisimulation_estimate)
average_target = tf.reduce_mean(bisimulation_target)
average_next_state_dists = tf.reduce_mean(
self.next_state_distances)
tf.summary.scalar('Training/loss', bisimulation_loss)
tf.summary.scalar('Training/AverageDistance', average_distance)
tf.summary.scalar('Training/AverageTargetDistance', average_target)
tf.summary.scalar('Training/AverageNextStateDistance',
average_next_state_dists)
tf.summary.scalar('Training/BisimHorizon', self.bisim_horizon_ph)
tf.summary.histogram('Training/OnlineDistance',
bisimulation_estimate)
tf.summary.histogram('Training/TargetDistance',
bisimulation_target)
self._train_bisim_op = self.bisim_optimizer.minimize(bisimulation_loss)
self._bisim_sync_op = self._build_sync_op(online_scope='OnlineBisim',
target_scope='TargetBisim')
def initNetworks(self):
net = tf.make_template('net', self.network)
if self.data_format == 'channels_first':
data_generator_entries = OrderedDict([
('image',
[self.image_channels] + list(reversed(self.image_size))),
('landmarks',
[self.num_landmarks] + list(reversed(self.heatmap_size)))
])
data_generator_entries_val = OrderedDict([
('image',
[self.image_channels] + list(reversed(self.image_size))),
('landmarks',
[self.num_landmarks] + list(reversed(self.heatmap_size)))
])
else:
raise NotImplementedError
self.train_queue = DataGenerator(self.dataset_train,
self.coord,
data_generator_entries,
batch_size=self.batch_size,
n_threads=8)
placeholders = self.train_queue.dequeue()
image = placeholders[0]
landmarks = placeholders[1]
prediction = net(image,
num_landmarks=self.num_landmarks,
is_training=True,
data_format=self.data_format)
self.loss_net = self.loss_function(landmarks, prediction)
with tf.control_dependencies(tf.get_collection(
tf.GraphKeys.UPDATE_OPS)):
if self.reg_constant > 0:
reg_losses = tf.get_collection(
tf.GraphKeys.REGULARIZATION_LOSSES)
self.loss_reg = self.reg_constant * tf.add_n(reg_losses)
self.loss = self.loss_net + self.loss_reg
else:
self.loss_reg = 0
self.loss = self.loss_net
self.train_losses = OrderedDict([('loss', self.loss_net),
('loss_reg', self.loss_reg)])
#self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate).minimize(self.loss)
self.optimizer = tf.train.MomentumOptimizer(
learning_rate=self.learning_rate, momentum=0.99,
use_nesterov=True).minimize(self.loss)
# build val graph
self.val_placeholders = tensorflow_train.utils.tensorflow_util.create_placeholders(
data_generator_entries_val, shape_prefix=[1])
self.image_val = self.val_placeholders['image']
self.landmarks_val = self.val_placeholders['landmarks']
self.prediction_val = net(self.image_val,
num_landmarks=self.num_landmarks,
is_training=False,
data_format=self.data_format)
# losses
self.loss_val = self.loss_function(self.landmarks_val,
self.prediction_val)
self.val_losses = OrderedDict([('loss', self.loss_val),
('loss_reg', self.loss_reg)])
def get_prob_discriminator(model_name, scope='prob_discriminator', **kwargs):
model_func = prob_discriminator_dict[model_name]
return tf.make_template(scope, model_func, **kwargs)
def main():
# initialize data loaders for train/test splits
if args.data_set == 'imagenet' and args.class_conditional:
raise("We currently don't have labels for the small imagenet data set")
if args.data_set == 'cifar':
import data.cifar10_data as cifar10_data
DataLoader = cifar10_data.DataLoader
elif args.data_set == 'imagenet':
import data.imagenet_data as imagenet_data
DataLoader = imagenet_data.DataLoader
else:
raise("unsupported dataset")
train_data = DataLoader(args.data_dir, 'train', args.batch_size * args.nr_gpu, rng=rng, shuffle=True, return_labels=args.class_conditional)
test_data = DataLoader(args.data_dir, 'test', args.batch_size * args.nr_gpu, shuffle=False, return_labels=args.class_conditional)
obs_shape = train_data.get_observation_size() # e.g. a tuple (32,32,3)
assert len(obs_shape) == 3, 'assumed right now'
# data place holders
x_init = tf.placeholder(tf.float32, shape=(args.init_batch_size,) + obs_shape)
xs = [tf.placeholder(tf.float32, shape=(args.batch_size, ) + obs_shape) for i in range(args.nr_gpu)]
# if the model is class-conditional we'll set up label placeholders + one-hot encodings 'h' to condition on
if args.class_conditional:
num_labels = train_data.get_num_labels()
y_init = tf.placeholder(tf.int32, shape=(args.init_batch_size,))
h_init = tf.one_hot(y_init, num_labels)
y_sample = np.split(np.mod(np.arange(args.batch_size*args.nr_gpu), num_labels), args.nr_gpu)
h_sample = [tf.one_hot(tf.Variable(y_sample[i], trainable=False), num_labels) for i in range(args.nr_gpu)]
ys = [tf.placeholder(tf.int32, shape=(args.batch_size,)) for i in range(args.nr_gpu)]
hs = [tf.one_hot(ys[i], num_labels) for i in range(args.nr_gpu)]
else:
h_init = None
h_sample = [None] * args.nr_gpu
hs = h_sample
# create the model
model_opt = { 'nr_resnet': args.nr_resnet, 'nr_filters': args.nr_filters, 'nr_logistic_mix': args.nr_logistic_mix, 'resnet_nonlinearity': args.resnet_nonlinearity}
model = tf.make_template('model', model_spec)
# run once for data dependent initialization of parameters
data_dependent_init = model(x_init, h_init, init=True, dropout_p=args.dropout_p, **model_opt)
# keep track of moving average
all_params = tf.trainable_variables()
ema = tf.train.ExponentialMovingAverage(decay=args.polyak_decay)
maintain_averages_op = tf.group(ema.apply(all_params))
ema_params = [ema.average(p) for p in all_params]
# get loss gradients over multiple GPUs + sampling
grads = []
loss_gen = []
loss_gen_test = []
new_x_gen = []
for i in range(args.nr_gpu):
with tf.device('/gpu:%d' % i):
if args.graph_cloning and i>0:
# already defined the graph once, use it again via template rather than redefining again
in_ = [xs[i]] + tf.global_variables()
res = gpu_template.apply(in_)
loss_train, loss_test, sx = res[:3]
grad = res[3:]
loss_gen.append(loss_train)
loss_gen_test.append(loss_test)
new_x_gen.append(sx)
grads.append(grad)
else:
# train
out = model(xs[i], hs[i], ema=None, dropout_p=args.dropout_p, **model_opt)
loss_gen.append(nn.discretized_mix_logistic_loss(tf.stop_gradient(xs[i]), out))
# gradients
grads.append(gradients(loss_gen[i], all_params))
# test
out = model(xs[i], hs[i], ema=ema, dropout_p=0., **model_opt)
loss_gen_test.append(nn.discretized_mix_logistic_loss(xs[i], out))
# sample
out = model(xs[i], h_sample[i], ema=ema, dropout_p=0, **model_opt)
new_x_gen.append(nn.sample_from_discretized_mix_logistic(out, args.nr_logistic_mix))
if args.graph_cloning:
in_ = [xs[0]] + tf.global_variables()
out_ = [loss_gen[0], loss_gen_test[0], new_x_gen[0]] + grads[0]
gpu_template = GraphTemplate(in_, outputs=out_)
# add losses and gradients together and get training updates
tf_lr = tf.placeholder(tf.float32, shape=[])
with tf.device('/gpu:0'):
for i in range(1,args.nr_gpu):
loss_gen[0] += loss_gen[i]
loss_gen_test[0] += loss_gen_test[i]
for j in range(len(grads[0])):
grads[0][j] += grads[i][j]
# training op
optimizer = tf.group(nn.adam_updates(all_params, grads[0], lr=tf_lr, mom1=0.95, mom2=0.9995), maintain_averages_op)
# convert loss to bits/dim
bits_per_dim = loss_gen[0]/(args.nr_gpu*np.log(2.)*np.prod(obs_shape)*args.batch_size)
bits_per_dim_test = loss_gen_test[0]/(args.nr_gpu*np.log(2.)*np.prod(obs_shape)*args.batch_size)
# sample from the model
def sample_from_model(sess):
x_gen = [np.zeros((args.batch_size,) + obs_shape, dtype=np.float32) for i in range(args.nr_gpu)]
for yi in range(obs_shape[0]):
for xi in range(obs_shape[1]):
new_x_gen_np = sess.run(new_x_gen, {xs[i]: x_gen[i] for i in range(args.nr_gpu)})
for i in range(args.nr_gpu):
x_gen[i][:,yi,xi,:] = new_x_gen_np[i][:,yi,xi,:]
return np.concatenate(x_gen, axis=0)
# turn numpy inputs into feed_dict for use with tensorflow
def make_feed_dict(data, init=False):
if type(data) is tuple:
x,y = data
else:
x = data
y = None
x = np.cast[np.float32]((x - 127.5) / 127.5) # input to pixelCNN is scaled from uint8 [0,255] to float in range [-1,1]
if init:
feed_dict = {x_init: x}
if y is not None:
feed_dict.update({y_init: y})
else:
x = np.split(x, args.nr_gpu)
feed_dict = {xs[i]: x[i] for i in range(args.nr_gpu)}
if y is not None:
y = np.split(y, args.nr_gpu)
feed_dict.update({ys[i]: y[i] for i in range(args.nr_gpu)})
return feed_dict
# //////////// perform training //////////////
if not os.path.exists(args.save_dir):
os.makedirs(args.save_dir)
print('starting training')
test_bpd = []
lr = args.learning_rate
saver = tf.train.Saver()
with tf.Session() as sess:
for epoch in range(args.max_epochs):
begin = time.time()
# init
if epoch == 0:
feed_dict = make_feed_dict(train_data.next(args.init_batch_size), init=True) # manually retrieve exactly init_batch_size examples
train_data.reset() # rewind the iterator back to 0 to do one full epoch
print('initializing the model...')
sess.run(tf.global_variables_initializer())
sess.run(data_dependent_init, feed_dict)
# train for one epoch
train_losses = []
counter = 0
for d in train_data:
counter+=1
feed_dict = make_feed_dict(d)
# forward/backward/update model on each gpu
lr *= args.lr_decay
feed_dict.update({ tf_lr: lr })
l,_ = sess.run([bits_per_dim, optimizer], feed_dict)
print(counter, l)
train_losses.append(l)
if counter>50:
if l>6.5:
assert False, "Test failed, expected loss 6.28537 at iteration 50"
else:
print("Test passed, loss %f (expected %f)"%(l, 6.28537))
sys.exit()
train_loss_gen = np.mean(train_losses)
# compute likelihood over test data
test_losses = []
for d in test_data:
feed_dict = make_feed_dict(d)
l = sess.run(bits_per_dim_test, feed_dict)
test_losses.append(l)
test_loss_gen = np.mean(test_losses)
test_bpd.append(test_loss_gen)
# log progress to console
print("Iteration %d, time = %ds, train bits_per_dim = %.4f, test bits_per_dim = %.4f" % (epoch, time.time()-begin, train_loss_gen, test_loss_gen))
sys.stdout.flush()
if epoch % args.save_interval == 0:
# generate samples from the model
sample_x = []
for i in range(args.num_samples):
sample_x.append(sample_from_model(sess))
sample_x = np.concatenate(sample_x,axis=0)
#img_tile = plotting.img_tile(sample_x[:100], aspect_ratio=1.0, border_color=1.0, stretch=True)
#img = plotting.plot_img(img_tile, title=args.data_set + ' samples')
#plotting.plt.savefig(os.path.join(args.save_dir,'%s_sample%d.png' % (args.data_set, epoch)))
#plotting.plt.close('all')
np.savez(os.path.join(args.save_dir,'%s_sample%d.npz' % (args.data_set, epoch)), sample_x)
# save params
saver.save(sess, args.save_dir + '/params_' + args.data_set + '.ckpt')
np.savez(args.save_dir + '/test_bpd_' + args.data_set + '.npz', test_bpd=np.array(test_bpd))
def __init__(self,
filters,
blocks,
use_nin,
components,
attn_heads,
use_ln,
with_affine=True,
use_final_nin=False,
init_scale=0.1,
nonlinearity=concat_elu):
self.components = components
self.with_affine = with_affine
self.scale_flow = Inverse(Sigmoid())
def f(x, init, ema, dropout_p, verbose, context):
if init and verbose:
with tf.variable_scope('debug'):
xmean, xvar = tf.nn.moments(x,
axes=list(
range(len(x.get_shape()))))
x = tf.Print(
x,
[
tf.shape(x), xmean,
tf.sqrt(xvar),
tf.reduce_min(x),
tf.reduce_max(x),
tf.reduce_any(tf.is_nan(x)),
tf.reduce_any(tf.is_inf(x))
],
message='{} (shape/mean/std/min/max/nan/inf) '.format(
self.template.variable_scope.name),
summarize=10,
)
B, H, W, C = x.shape.as_list()
pos_emb = to_default_floatx(
get_var(
'pos_emb',
ema=ema,
shape=[H, W, filters],
initializer=tf.random_normal_initializer(stddev=0.01),
))
x = conv2d(x, name='c1', num_units=filters, init=init, ema=ema)
for i_block in range(blocks):
with tf.variable_scope('block{}'.format(i_block)):
x = gated_resnet(x,
name='conv',
a=context,
use_nin=use_nin,
init=init,
ema=ema,
dropout_p=dropout_p)
if use_ln:
x = norm(x, name='ln1', ema=ema)
x = attn(x,
name='attn',
pos_emb=pos_emb,
heads=attn_heads,
init=init,
ema=ema,
dropout_p=dropout_p)
if use_ln:
x = norm(x, name='ln2', ema=ema)
assert x.shape == [B, H, W, filters]
x = nonlinearity(x)
x = (nin if use_final_nin else conv2d)(x,
name='c2',
num_units=C *
(2 + 3 * components),
init_scale=init_scale,
init=init,
ema=ema)
assert x.shape == [B, H, W, C * (2 + 3 * components)]
x = tf.reshape(x, [B, H, W, C, 2 + 3 * components])
x = at_least_float32(x) # do mix-logistics stuff in float32
s, t = tf.tanh(x[:, :, :, :, 0]), x[:, :, :, :, 1]
ml_logits, ml_means, ml_logscales = tf.split(x[:, :, :, :, 2:],
3,
axis=4)
ml_logscales = tf.maximum(ml_logscales, -7.)
assert s.shape == t.shape == [B, H, W, C]
assert ml_logits.shape == ml_means.shape == ml_logscales.shape == [
B, H, W, C, components
]
return s, t, ml_logits, ml_means, ml_logscales
self.template = tf.make_template(self.__class__.__name__, f)
# This allows us to evaluate the probability of an image under the distribution, not just individual pixels.
# See for below: /Users/coder352/github/jImage/Machine_Learning/VAE_encoder-prior-decoder.png
def plot_codes(ax, codes, labels):
ax.scatter(codes[:, 0], codes[:, 1], s=2, c=labels, alpha=0.1)
ax.set_aspect('equal')
ax.set_xlim(codes.min() - .1, codes.max() + .1)
ax.set_ylim(codes.min() - .1, codes.max() + .1)
ax.tick_params(
axis='both', which='both', left='off', bottom='off',
labelleft='off', labelbottom='off')
def plot_samples(ax, samples):
for index, sample in enumerate(samples):
ax[index].imshow(sample, cmap='gray')
ax[index].axis('off')
make_encoder = tf.make_template('encoder', make_encoder) # since make_prior doesn't have the tf.layers functions, so it doesn't have Variable
make_decoder = tf.make_template('decoder', make_decoder)
# Define the model.
prior = make_prior(code_size=2) # because there has height and width attributes for a grey image
posterior = make_encoder(data, code_size=2)
code = posterior.sample()
# Define the loss.
likelihood = make_decoder(code, [28, 28]).log_prob(data)
divergence = tfd.kl_divergence(posterior, prior)
elbo = tf.reduce_mean(likelihood - divergence)
optimize = tf.train.AdamOptimizer(0.001).minimize(-elbo)
samples = make_decoder(prior.sample(10), [28, 28]).mean()
models = [
MLPRegressor(counters={}, user_mode=args.user_mode)
for i in range(args.nr_model)
]
model_opt = {
"mlp": mlp,
"obs_shape": [1],
"alpha": 0.01,
"nonlinearity": tf.nn.relu,
"bn": False,
"kernel_initializer": tf.contrib.layers.xavier_initializer(uniform=False),
"kernel_regularizer": None,
}
model = tf.make_template('model', MLPRegressor.construct)
for i in range(args.nr_model):
with tf.device('/' + args.device_type + ':%d' % (i % args.nr_gpu)):
model(models[i], **model_opt)
save_dir = "/data/ziz/jxu/maml/test-{0}".format(args.dataset_name)
learner = MAMLLearner(session=None,
parallel_models=models,
optimize_op=None,
train_set=train_set,
eval_set=val_set,
variables=tf.trainable_variables(),
lr=args.learning_rate,
device_type=args.device_type,
save_dir=save_dir)
@property
@templatemethod("trial")
def trial(self):
print("call trial here")
w = tf.get_variable('w', [])
return tf.reduce_sum(self._x) * w
def my_trail(x, share_variable_name):
var1 = tf.get_variable(share_variable_name, shape=[])
return tf.reduce_sum(x) * var1
template_my = tf.make_template("template_my",
my_trail,
share_variable_name="my_v")
def test_model(source_path, target_path, vocab_path):
tf.logging.set_verbosity(tf.logging.INFO)
batch_size = 2
# Build model graph
mode = tf.contrib.learn.ModeKeys.TRAIN
params_ = AttentionSeq2Seq.default_params().copy()
params_.update({
"vocab_source": vocab_path,
"vocab_target": vocab_path,
})
def main():
"""
Simple demonstration of how you can put a GP on top of a NN and train the whole system end-to-end in GPflow-1.0.
Note
that in the new GPflow there are new features that we do not take advantage of here but could be used to make
the whole example cleaner. For example you may want to use a gpflow.train Optimiser as this will take care of
passing in the GP model feed dict for you as well as initially initialising the optimisers TF variables.
You could also choose to tell the gpmodel to initialise the NN variables by subclassing SVGP and overriding the
appropriate variable initialisation method.
"""
# ## We load in the MNIST data. We will create a validation set but will not use it in this simple example.
dataset = np.loadtxt("test", delimiter=",")
x_test = dataset[:,0:72]
y_test = dataset[:,72].reshape(-1,1)
#print(x_test[20])
dataset = np.loadtxt("dev", delimiter=",")
x_valid = dataset[:,0:72]
y_valid = dataset[:,72].reshape(-1,1)
#resampleFile()
dataset = np.loadtxt("train.revised", delimiter=",")
x_train = dataset[:,0:72]
y_train = dataset[:,72].reshape(-1,1)
x_train_root = x_train
x_valid_root = x_valid
x_train, x_test, x_valid = standardize_data(copy.deepcopy(x_train_root), x_test, copy.deepcopy(x_valid_root))
#rng = np.random.RandomState(100)
#train_permute = rng.permutation(x_train.shape[0])
#x_train, y_train = x_train[train_permute, :], y_train[train_permute, :]
# ## We set up a TensorFlow Graph and a Session linked to this.
tf_graph = tf.Graph()
tf_session = tf.Session(graph=tf_graph)
# ## We have some settings for the model and its training which we will set up below.
num_h = 17
num_classes = 2 #could be improved here
num_inducing = 100
minibatch_size = 250
# ## We set up the NN part of the GP kernel. This needs to be put on the same graph
with tf_graph.as_default():
phs = DataPlaceholders()
nn_base = tf.make_template("sconvnet_kernel", make_feedforward_nn, end_h=num_h) # end h is the number of hidden
# units at the end
h = nn_base(phs.data, phs.keep_prob)
h = tf.cast(h, gpf.settings.tf_float)
nn_vars = tf.global_variables() # only nn variables exist up to now.
tf_session.run(tf.variables_initializer(nn_vars))
# ## We now set up the GP part. Instead of the usual X data it will get the data after being processed by the NN.
with gpf.defer_build():
kernel = gpf.kernels.Matern52(num_h)
likelihood = gpf.likelihoods.MultiClass(num_classes)
gp_model = gpf.models.SVGP(h, phs.label, kernel, likelihood, np.ones((num_inducing, num_h), gpf.settings.np_float),
num_latent=num_classes, whiten=False, minibatch_size=None, num_data=x_train.shape[0])
# ^ so we say minibatch size is None to make sure we get DataHolder rather than minibatch data holder, which
# does not allow us to give in tensors. But we will handle all our minibatching outside.
gp_model.compile(tf_session)
# ## The initial lengthscales and inducing point locations are likely very bad. So we use heuristics for good
# initial starting points and reset them at these values.
gp_model.feature.Z.assign(suggest_good_intial_inducing_points(phs, np.unique(x_train, axis=0)[:5000, :], h, tf_session, num_inducing), tf_session)
gp_model.kern.lengthscales.assign(suggest_sensible_lengthscale(phs, np.unique(x_train, axis=0)[:5000, :], h, tf_session)
+ np.zeros_like(gp_model.kern.lengthscales.read_value()), tf_session)
# ^ note that this assign should reapply the transform for us :). The zeros ND array exists to make sure
# the lengthscales are the correct shape via broadcasting
# ## We create ops to measure the predictive log likelihood and the accuracy.
with tf_graph.as_default():
log_likelihood_predict = gp_model.likelihood.predict_density(*gp_model._build_predict(h), phs.label)
outputs = tf.argmax(gp_model.likelihood.predict_mean_and_var(*gp_model._build_predict(h))[0], axis=1, output_type=tf.int32)
accuracy = tf.cast(tf.equal(outputs, tf.squeeze(phs.label)), tf.float32)
avg_acc = tf.reduce_mean(accuracy)
avg_ll = tf.reduce_mean(log_likelihood_predict)
# ## we now create an optimiser and initialise its variables. Note that you could use a GPflow optimiser here
# and this would now be done for you.
all_vars_up_to_trainer = tf.global_variables()
optimiser = tf.train.AdamOptimizer(1e-4)
print(tf.global_variables())
minimise = optimiser.minimize(gp_model.objective) # this should pick up all Trainable variables.
adam_vars = list(set(tf.global_variables()) - set(all_vars_up_to_trainer))
neg_gp_opj = -gp_model.objective
tf_session.run(tf.variables_initializer(adam_vars))
saver = tf.train.Saver()
# ## We now go through a training loop where we optimise the NN and GP. we will print out the test results at
# regular intervals.
results = []
print("starting")
SEED = 448
np.random.RandomState(SEED)
for i in range(50): #100 epochs
if i>0:
#load truoc
#print_tensors_in_checkpoint_file(file_name="model_at_epoch"+str(i-1)+".ckpt", tensor_name='', all_tensors=False)
saver.restore(tf_session, "model_at_epoch"+str(i-1)+".ckpt")
fd = gp_model.feeds or {}
fd.update({phs.keep_prob: 1.0, phs.data: x_valid,
phs.label: y_valid})
accuracy_evald, log_like_evald, outputs_evald = tf_session.run([avg_acc, avg_ll, outputs], feed_dict=fd)
print("Epoch {}: \Dev set LL {}, Acc {}, Outputs {}".format(i, log_like_evald,
accuracy_evald, outputs_evald))
outputs_evald = outputs_evald.reshape(-1, 1)
#print(outputs_evald)
print("Result from the previous epoch on dev:")
compute_scores(y_valid, outputs_evald)
fd = gp_model.feeds or {}
fd.update({phs.keep_prob: 1.0, phs.data: x_test,
phs.label: y_test})
accuracy_evald, log_like_evald, outputs_evald = tf_session.run([avg_acc, avg_ll, outputs], feed_dict=fd)
print("Epoch {}: \nTest set LL {}, Acc {}, Outputs {}".format(i, log_like_evald,
accuracy_evald, outputs_evald))
outputs_evald = outputs_evald.reshape(-1, 1)
#print(outputs_evald)
print("Result from the previous epoch on test:")
compute_scores(y_test, outputs_evald)
shuffle = np.arange(y_train.size)
np.random.shuffle(shuffle)
print(shuffle)
x_train_shuffle = x_train[shuffle]
y_train_shuffle = y_train[shuffle]
data_indx = 0
while data_indx<y_train.size:
lastIndex = data_indx + minibatch_size
if lastIndex>=y_train.size:
lastIndex = y_train.size
indx_array = np.mod(np.arange(data_indx, lastIndex), x_train_shuffle.shape[0])
#print("array", indx_array)
data_indx += minibatch_size
#print(data_indx)fz
fd = gp_model.feeds or {}
fd.update({
phs.keep_prob: 1.0,
phs.data: x_train_shuffle[indx_array],
phs.label: y_train_shuffle[indx_array]
})
_, loss_evd = tf_session.run([minimise, neg_gp_opj], feed_dict=fd)
# Print progress every 1 epoch.
save_path = saver.save(tf_session, "./model_at_epoch"+str(i)+".ckpt")
print("Done!")
def main(args):
import os
import sys
import time
import json
from mpi4py import MPI
import numpy as np
import tensorflow as tf
from tqdm import trange
import pixel_cnn_pp.nn as nn
import pixel_cnn_pp.plotting as plotting
from pixel_cnn_pp import model as pxpp_models
import data.cifar10_data as cifar10_data
import data.imagenet_data as imagenet_data
import tf_utils as tfu
comm = MPI.COMM_WORLD
num_tasks, task_id = comm.Get_size(), comm.Get_rank()
save_dir = args.save_dir
if task_id == 0:
os.makedirs(save_dir)
f_log = open(os.path.join(save_dir, 'print.log'), 'w')
def lprint(*a, **kw):
if task_id == 0:
print(*a, **kw)
print(*a, **kw, file=f_log)
lprint('input args:\n', json.dumps(vars(args), indent=4,
separators=(',', ':'))) # pretty print args
# -----------------------------------------------------------------------------
# fix random seed for reproducibility
rng = np.random.RandomState(args.seed + task_id)
tf.set_random_seed(args.seed + task_id)
# initialize data loaders for train/test splits
if args.data_set == 'imagenet' and args.class_conditional:
raise("We currently don't have labels for the small imagenet data set")
DataLoader = {'cifar': cifar10_data.DataLoader,
'imagenet': imagenet_data.DataLoader}[args.data_set]
train_data = DataLoader(args.data_dir, 'train', args.batch_size,
rng=rng, shuffle=True, return_labels=args.class_conditional)
test_data = DataLoader(args.data_dir, 'test', args.batch_size,
shuffle=False, return_labels=args.class_conditional)
obs_shape = train_data.get_observation_size() # e.g. a tuple (32,32,3)
assert len(obs_shape) == 3, 'assumed right now'
if args.nr_gpu is None:
from tensorflow.python.client import device_lib
args.nr_gpu = len([d for d in device_lib.list_local_devices()
if d.device_type == 'GPU'])
# data place holders
x_init = tf.placeholder(tf.float32,
shape=(args.init_batch_size,) + obs_shape)
xs = [tf.placeholder(tf.float32, shape=(args.batch_size, ) + obs_shape)
for _ in range(args.nr_gpu)]
def _get_batch(is_training):
if is_training:
x = train_data.__next__(args.batch_size)
else:
x = test_data.__next__(args.batch_size)
x = np.cast[np.float32]((x - 127.5) / 127.5)
return dict(x=x)
batch_def = dict(x=tfu.vdef(args.batch_size, obs_shape))
qr = tfu.Struct(
train=tfu.PyfuncRunner(batch_def, 64, 8, True,
_get_batch, is_training=True),
test=tfu.PyfuncRunner(batch_def, 64, 8, True,
_get_batch, is_training=False),
)
tf.add_to_collection(tf.GraphKeys.QUEUE_RUNNERS, qr.train)
tf.add_to_collection(tf.GraphKeys.QUEUE_RUNNERS, qr.test)
if args.nr_gpu is None:
from tensorflow.python.client import device_lib
args.nr_gpu = len([d for d in device_lib.list_local_devices()
if d.device_type == 'GPU'])
sess = tfu.Session(allow_soft_placement=True).__enter__()
# if the model is class-conditional we'll set up label placeholders +
# one-hot encodings 'h' to condition on
if args.class_conditional:
raise NotImplementedError
num_labels = train_data.get_num_labels()
y_init = tf.placeholder(tf.int32, shape=(args.init_batch_size,))
h_init = tf.one_hot(y_init, num_labels)
y_sample = np.split(
np.mod(np.arange(args.batch_size), num_labels), args.nr_gpu)
h_sample = [tf.one_hot(tf.Variable(y_sample[i], trainable=False), num_labels)
for i in range(args.nr_gpu)]
ys = [tf.placeholder(tf.int32, shape=(args.batch_size,))
for i in range(args.nr_gpu)]
hs = [tf.one_hot(ys[i], num_labels) for i in range(args.nr_gpu)]
else:
h_init = None
h_sample = [None] * args.nr_gpu
hs = h_sample
# create the model
model_opt = {'nr_resnet': args.nr_resnet, 'nr_filters': args.nr_filters,
'nr_logistic_mix': args.nr_logistic_mix, 'resnet_nonlinearity': args.resnet_nonlinearity}
model = tf.make_template('model', getattr(pxpp_models, args.model + "_spec"))
# run once for data dependent initialization of parameters
with tf.device('/gpu:0'):
gen_par = model(x_init, h_init, init=True,
dropout_p=args.dropout_p, **model_opt)
# keep track of moving average
all_params = tf.trainable_variables()
lprint('# of Parameters', sum(np.prod(p.get_shape().as_list())
for p in all_params))
ema = tf.train.ExponentialMovingAverage(decay=args.polyak_decay)
maintain_averages_op = tf.group(ema.apply(all_params))
loss_gen, loss_gen_test, grads = [], [], []
for i in range(args.nr_gpu):
with tf.device('/gpu:%d' % i):
x = qr.train.batch().x
gen_par = model(x, hs[i], ema=None,
dropout_p=args.dropout_p, **model_opt)
if isinstance(gen_par, tuple) and len(gen_par) == 3:
loss_gen.append(nn.discretized_mix_logistic_loss_per_chn(x, *gen_par))
else:
loss_gen.append(nn.discretized_mix_logistic_loss(x, gen_par))
grads.append(tf.gradients(loss_gen[i], all_params))
x = qr.test.batch().x
gen_par = model(x, hs[i], ema=ema, dropout_p=0., **model_opt)
if isinstance(gen_par, tuple) and len(gen_par) == 3:
loss_gen_test.append(
nn.discretized_mix_logistic_loss_per_chn(x, *gen_par))
else:
loss_gen_test.append(nn.discretized_mix_logistic_loss(x, gen_par))
# add losses and gradients together and get training updates
tf_lr = tf.placeholder(tf.float32, shape=[])
with tf.device('/gpu:0'):
for i in range(1, args.nr_gpu):
loss_gen[0] += loss_gen[i]
loss_gen_test[0] += loss_gen_test[i]
for j in range(len(grads[0])):
grads[0][j] += grads[i][j]
if num_tasks > 1:
lprint('creating mpi optimizer')
# If we have multiple mpi processes, average across them.
flat_grad = tf.concat([tf.reshape(g, (-1,)) for g in grads[0]], axis=0)
shapes = [g.shape.as_list() for g in grads[0]]
sizes = [int(np.prod(s)) for s in shapes]
buf = np.zeros(sum(sizes), np.float32)
def _gather_grads(my_flat_grad):
comm.Allreduce(my_flat_grad, buf, op=MPI.SUM)
np.divide(buf, float(num_tasks), out=buf)
return buf
avg_flat_grad = tf.py_func(_gather_grads, [flat_grad], tf.float32)
avg_flat_grad.set_shape(flat_grad.shape)
avg_grads = tf.split(avg_flat_grad, sizes, axis=0)
grads[0] = [tf.reshape(g, v.shape) for g, v in zip(avg_grads, grads[0])]
# training op
optimizer = tf.group(nn.adam_updates(
all_params, grads[0], lr=tf_lr, mom1=0.95, mom2=0.9995, eps=1e-6), maintain_averages_op)
# convert loss to bits/dim
total_gpus = sum(comm.allgather(args.nr_gpu))
lprint('using %d gpus across %d machines' % (total_gpus, num_tasks))
norm_const = np.log(2.) * np.prod(obs_shape) * args.batch_size
norm_const *= total_gpus / num_tasks
bits_per_dim = loss_gen[0] / norm_const
bits_per_dim_test = loss_gen_test[0] / norm_const
bits_per_dim = tf.check_numerics(bits_per_dim, 'train loss is nan')
bits_per_dim_test = tf.check_numerics(bits_per_dim_test, 'test loss is nan')
new_x_gen = []
for i in range(args.nr_gpu):
with tf.device('/gpu:%d' % i):
gen_par = model(xs[i], hs[i], ema=ema, dropout_p=0, **model_opt)
new_x_gen.append(
nn.sample_from_discretized_mix_logistic(gen_par, args.nr_logistic_mix))
def sample_from_model(sess, n_samples=args.nr_gpu * args.batch_size):
sample_x = np.zeros((0,) + obs_shape, dtype=np.float32)
while len(sample_x) < n_samples:
x_gen = [np.zeros((args.batch_size,) + obs_shape, dtype=np.float32)
for i in range(args.nr_gpu)]
for yi in range(obs_shape[0]):
for xi in range(obs_shape[1]):
new_x_gen_np = sess.run(new_x_gen,
{xs[i]: x_gen[i] for i in range(args.nr_gpu)})
for i in range(args.nr_gpu):
x_gen[i][:, yi, xi, :] = new_x_gen_np[i][:, yi, xi, :]
sample_x = np.concatenate([sample_x] + x_gen, axis=0)
img_tile = plotting.img_tile(
sample_x[:int(np.floor(np.sqrt(n_samples))**2)],
aspect_ratio=1.0, border_color=1.0, stretch=True)
img = plotting.plot_img(img_tile, title=args.data_set + ' samples')
plotting.plt.savefig(
os.path.join(save_dir, '%s_samples.png' % args.data_set))
np.save(os.path.join(save_dir, '%s_samples.npy' % args.data_set), sample_x)
plotting.plt.close('all')
# init & save
initializer = tf.global_variables_initializer()
saver = tf.train.Saver()
# turn numpy inputs into feed_dict for use with tensorflow
def make_feed_dict(data, init=False):
if type(data) is tuple:
x, y = data
else:
x = data
y = None
# input to pixelCNN is scaled from uint8 [0,255] to float in range [-1,1]
x = np.cast[np.float32]((x - 127.5) / 127.5)
if init:
feed_dict = {x_init: x}
if y is not None:
feed_dict.update({y_init: y})
else:
x = np.split(x, args.nr_gpu)
feed_dict = {xs[i]: x[i] for i in range(args.nr_gpu)}
if y is not None:
y = np.split(y, args.nr_gpu)
feed_dict.update({ys[i]: y[i] for i in range(args.nr_gpu)})
return feed_dict
# //////////// perform training //////////////
lprint('dataset size: %d' % len(train_data.data))
test_bpd = []
lr = args.learning_rate
# manually retrieve exactly init_batch_size examples
feed_dict = make_feed_dict(train_data.next(args.init_batch_size), init=True)
train_data.reset() # rewind the iterator back to 0 to do one full epoch
lprint('initializing the model...')
sess.run(initializer, feed_dict)
if args.load_params:
# ckpt_file = save_dir + '/params_' + args.data_set + '.ckpt'
ckpt_file = args.load_params
lprint('restoring parameters from', ckpt_file)
saver.restore(sess, ckpt_file)
# Sync params before starting.
my_vals = sess.run(all_params)
vals = [np.zeros_like(v) for v in my_vals]
[comm.Allreduce(mv, v, op=MPI.SUM) for mv, v in zip(my_vals, vals)]
assign_ops = [var.assign(val / num_tasks)
for var, val in zip(all_params, vals)]
sess.run(assign_ops)
coord = tfu.start_queue_runners(sess)
batch_size = args.batch_size * total_gpus
iters_per_train_epoch = len(train_data.data) // batch_size
iters_per_test_epoch = len(test_data.data) // batch_size
lprint('starting training')
for epoch in range(args.max_epochs):
begin = time.time()
# train for one epoch
train_losses = []
ti = trange(iters_per_train_epoch)
for itr in ti:
if coord.should_stop():
tfu.stop_queue_runners(coord)
# forward/backward/update model on each gpu
lr *= args.lr_decay
l, _ = sess.run([bits_per_dim, optimizer], {tf_lr: lr})
train_losses.append(l)
ti.set_postfix(loss=l, lr=lr)
train_loss_gen = np.mean(train_losses)
# compute likelihood over test data
test_losses = []
for itr in trange(iters_per_test_epoch):
if coord.should_stop():
tfu.stop_queue_runners(coord)
l = sess.run(bits_per_dim_test)
test_losses.append(l)
test_loss_gen = np.mean(test_losses)
test_bpd.append(test_loss_gen)
# log progress to console
stats = dict(epoch=epoch, time=time.time() - begin, lr=lr,
train_bpd=train_loss_gen,
test_bpd=test_loss_gen)
all_stats = comm.gather(stats)
if task_id == 0:
lprint('-' * 16)
for k in stats:
lprint('%s:\t%s' % (k, np.mean([s[k] for s in all_stats])))
if epoch % args.save_interval == 0:
path = os.path.join(save_dir, str(epoch))
os.makedirs(path, exist_ok=True)
saver.save(sess, os.path.join(path, 'params_%s.ckpt' % args.data_set))
sample_from_model(sess)
def masked_autoregressive_default_template(hidden_layers,
shift_only=False,
activation=tf.nn.relu,
log_scale_min_clip=-5.,
log_scale_max_clip=3.,
log_scale_clip_gradient=False,
name=None,
*args, # pylint: disable=keyword-arg-before-vararg
**kwargs):
"""Build the Masked Autoregressive Density Estimator (Germain et al., 2015).
This will be wrapped in a make_template to ensure the variables are only
created once. It takes the input and returns the `loc` ("mu" in [Germain et
al. (2015)][1]) and `log_scale` ("alpha" in [Germain et al. (2015)][1]) from
the MADE network.
Warning: This function uses `masked_dense` to create randomly initialized
`tf.Variables`. It is presumed that these will be fit, just as you would any
other neural architecture which uses `tf.layers.dense`.
#### About Hidden Layers
Each element of `hidden_layers` should be greater than the `input_depth`
(i.e., `input_depth = tf.shape(input)[-1]` where `input` is the input to the
neural network). This is necessary to ensure the autoregressivity property.
#### About Clipping
This function also optionally clips the `log_scale` (but possibly not its
gradient). This is useful because if `log_scale` is too small/large it might
underflow/overflow making it impossible for the `MaskedAutoregressiveFlow`
bijector to implement a bijection. Additionally, the `log_scale_clip_gradient`
`bool` indicates whether the gradient should also be clipped. The default does
not clip the gradient; this is useful because it still provides gradient
information (for fitting) yet solves the numerical stability problem. I.e.,
`log_scale_clip_gradient = False` means
`grad[exp(clip(x))] = grad[x] exp(clip(x))` rather than the usual
`grad[clip(x)] exp(clip(x))`.
Args:
hidden_layers: Python `list`-like of non-negative integer, scalars
indicating the number of units in each hidden layer. Default: `[512, 512].
shift_only: Python `bool` indicating if only the `shift` term shall be
computed. Default: `False`.
activation: Activation function (callable). Explicitly setting to `None`
implies a linear activation.
log_scale_min_clip: `float`-like scalar `Tensor`, or a `Tensor` with the
same shape as `log_scale`. The minimum value to clip by. Default: -5.
log_scale_max_clip: `float`-like scalar `Tensor`, or a `Tensor` with the
same shape as `log_scale`. The maximum value to clip by. Default: 3.
log_scale_clip_gradient: Python `bool` indicating that the gradient of
`tf.clip_by_value` should be preserved. Default: `False`.
name: A name for ops managed by this function. Default:
"masked_autoregressive_default_template".
*args: `tf.layers.dense` arguments.
**kwargs: `tf.layers.dense` keyword arguments.
Returns:
shift: `Float`-like `Tensor` of shift terms (the "mu" in
[Germain et al. (2015)][1]).
log_scale: `Float`-like `Tensor` of log(scale) terms (the "alpha" in
[Germain et al. (2015)][1]).
Raises:
NotImplementedError: if rightmost dimension of `inputs` is unknown prior to
graph execution.
#### References
[1]: Mathieu Germain, Karol Gregor, Iain Murray, and Hugo Larochelle. MADE:
Masked Autoencoder for Distribution Estimation. In _International
Conference on Machine Learning_, 2015. https://arxiv.org/abs/1502.03509
"""
name = name or "masked_autoregressive_default_template"
with tf.name_scope(name, values=[log_scale_min_clip, log_scale_max_clip]):
def _fn(x):
"""MADE parameterized via `masked_autoregressive_default_template`."""
# TODO(b/67594795): Better support of dynamic shape.
input_depth = x.shape.with_rank_at_least(1)[-1].value
if input_depth is None:
raise NotImplementedError(
"Rightmost dimension must be known prior to graph execution.")
input_shape = (
np.int32(x.shape.as_list())
if x.shape.is_fully_defined() else tf.shape(x))
for i, units in enumerate(hidden_layers):
x = masked_dense(
inputs=x,
units=units,
num_blocks=input_depth,
exclusive=True if i == 0 else False,
activation=activation,
*args, # pylint: disable=keyword-arg-before-vararg
**kwargs)
x = masked_dense(
inputs=x,
units=(1 if shift_only else 2) * input_depth,
num_blocks=input_depth,
activation=None,
*args, # pylint: disable=keyword-arg-before-vararg
**kwargs)
if shift_only:
x = tf.reshape(x, shape=input_shape)
return x, None
x = tf.reshape(x, shape=tf.concat([input_shape, [2]], axis=0))
shift, log_scale = tf.unstack(x, num=2, axis=-1)
which_clip = (
tf.clip_by_value
if log_scale_clip_gradient else _clip_by_value_preserve_grad)
log_scale = which_clip(log_scale, log_scale_min_clip, log_scale_max_clip)
return shift, log_scale
return tf.make_template(name, _fn)
def func_wrapper(*args, **kwargs):
"""Inner wrapper function"""
templated_func = tf.make_template(name_, func)
return templated_func(*args, **kwargs)
def editor_train(base_words,
extended_base_words,
output_words,
extended_output_words,
source_words,
target_words,
insert_words,
delete_words,
oov,
vocab_size,
hidden_dim,
agenda_dim,
edit_dim,
micro_edit_ev_dim,
num_heads,
num_encoder_layers,
num_decoder_layers,
attn_dim,
beam_width,
ctx_hidden_dim,
ctx_hidden_layer,
wa_hidden_dim,
wa_hidden_layer,
meve_hidden_dim,
meve_hidden_layers,
max_sent_length,
dropout_keep,
lamb_reg,
norm_eps,
norm_max,
kill_edit,
draw_edit,
swap_memory,
use_beam_decoder=False,
use_dropout=False,
no_insert_delete_attn=False,
enable_vae=True):
batch_size = tf.shape(source_words)[0]
# [batch]
base_len = seq.length_pre_embedding(base_words)
output_len = seq.length_pre_embedding(extended_output_words)
src_len = seq.length_pre_embedding(source_words)
tgt_len = seq.length_pre_embedding(target_words)
iw_len = seq.length_pre_embedding(insert_words)
dw_len = seq.length_pre_embedding(delete_words)
# variable of shape [vocab_size, embed_dim]
embeddings = vocab.get_embeddings()
# [batch x max_len x embed_dim]
base_word_embeds = vocab.embed_tokens(base_words)
src_word_embeds = vocab.embed_tokens(source_words)
tgt_word_embeds = vocab.embed_tokens(target_words)
insert_word_embeds = vocab.embed_tokens(insert_words)
delete_word_embeds = vocab.embed_tokens(delete_words)
sent_encoder = tf.make_template('sent_encoder',
encoder.source_sent_encoder,
hidden_dim=hidden_dim,
num_layer=num_encoder_layers,
swap_memory=swap_memory,
use_dropout=use_dropout,
dropout_keep=dropout_keep)
# [batch x max_len x rnn_out_dim], [batch x rnn_out_dim]
base_sent_hidden_states, base_sent_embed = sent_encoder(
base_word_embeds, base_len)
assert kill_edit == False and draw_edit == False
# [batch x edit_dim]
if kill_edit:
edit_vector = tf.zeros(shape=(batch_size, edit_dim))
else:
if draw_edit:
edit_vector = random_noise_encoder(batch_size, edit_dim, norm_max)
else:
edit_vector, wa_inserted, wa_deleted = attn_encoder(
src_word_embeds,
tgt_word_embeds,
insert_word_embeds,
delete_word_embeds,
src_len,
tgt_len,
iw_len,
dw_len,
ctx_hidden_dim,
ctx_hidden_layer,
wa_hidden_dim,
wa_hidden_layer,
meve_hidden_dim,
meve_hidden_layers,
edit_dim,
micro_edit_ev_dim,
num_heads,
lamb_reg,
norm_eps,
norm_max,
sent_encoder,
use_dropout=use_dropout,
dropout_keep=dropout_keep,
swap_memory=swap_memory,
enable_vae=enable_vae)
# [batch x agenda_dim]
base_agenda = agn.linear(base_sent_embed, edit_vector, agenda_dim)
train_dec_inp, train_dec_inp_len, \
train_dec_out, train_dec_out_len = prepare_decoder_input_output(output_words, extended_output_words, output_len)
train_dec_inp_extended = prepare_decoder_inputs(extended_output_words,
tf.cast(-1, tf.int64))
train_decoder = decoder.train_decoder(
base_agenda,
embeddings,
extended_base_words,
oov,
train_dec_inp,
train_dec_inp_extended,
base_sent_hidden_states,
wa_inserted,
wa_deleted,
train_dec_inp_len,
base_len,
src_len,
tgt_len,
vocab_size,
attn_dim,
hidden_dim,
num_decoder_layers,
swap_memory,
enable_dropout=use_dropout,
dropout_keep=dropout_keep,
no_insert_delete_attn=no_insert_delete_attn)
if use_beam_decoder:
infr_decoder = decoder.beam_eval_decoder(
base_agenda,
embeddings,
extended_base_words,
oov,
vocab.get_token_id(vocab.START_TOKEN),
vocab.get_token_id(vocab.STOP_TOKEN),
base_sent_hidden_states,
wa_inserted,
wa_deleted,
base_len,
src_len,
tgt_len,
vocab_size,
attn_dim,
hidden_dim,
num_decoder_layers,
max_sent_length,
beam_width,
swap_memory,
enable_dropout=use_dropout,
dropout_keep=dropout_keep,
no_insert_delete_attn=no_insert_delete_attn)
else:
infr_decoder = decoder.greedy_eval_decoder(
base_agenda,
embeddings,
extended_base_words,
oov,
vocab.get_token_id(vocab.START_TOKEN),
vocab.get_token_id(vocab.STOP_TOKEN),
base_sent_hidden_states,
wa_inserted,
wa_deleted,
base_len,
src_len,
tgt_len,
vocab_size,
attn_dim,
hidden_dim,
num_decoder_layers,
max_sent_length,
swap_memory,
enable_dropout=use_dropout,
dropout_keep=dropout_keep,
no_insert_delete_attn=no_insert_delete_attn)
return train_decoder, infr_decoder, train_dec_out, train_dec_out_len
train_input_single = tf.placeholder(tf.float32,
shape=(IMG_SIZE[0], IMG_SIZE[1],
1))
train_gt_single = tf.placeholder(tf.float32,
shape=(IMG_SIZE[0], IMG_SIZE[1], 1))
q = tf.FIFOQueue(
10000, [tf.float32, tf.float32],
[[IMG_SIZE[0], IMG_SIZE[1], 1], [IMG_SIZE[0], IMG_SIZE[1], 1]])
enqueue_op = q.enqueue([train_input_single, train_gt_single])
train_input, train_gt = q.dequeue_many(BATCH_SIZE)
### WITH ASYNCHRONOUS DATA LOADING ###
shared_model = tf.make_template('shared_model', model)
# train_output, weights = model(train_input)
train_output, weights = shared_model(train_input)
# this is the output of the VDSR.modle, now we add a sigma regulation to the loss function
# at the same time, we need give the truth to the weight_filter()
# pre_edge = Weight_filter.filter(train_output)
# gt_edge = Weight_filter.filter(train_gt)
loss = tf.reduce_sum(tf.nn.l2_loss(tf.subtract(train_output, train_gt)))
# + tf.reduce_sum(
# tf.nn.l2_loss(tf.subtract(pre_edge, gt_edge)))
for w in weights:
loss += tf.nn.l2_loss(w) * 1e-4
tf.summary.scalar("loss", loss)
global_step = tf.Variable(0, trainable=False)
learning_rate = tf.train.exponential_decay(BASE_LR,
global_step * BATCH_SIZE,
def __init__(self,
mem_slots,
head_size,
num_heads=1,
num_blocks=1,
forget_bias=1.0,
input_bias=0.0,
gate_style='unit',
attention_mlp_layers=2,
key_size=None,
name='relational_memory'):
"""Constructs a `RelationalMemory` object.
Args:
mem_slots: The total number of memory slots to use.
head_size: The size of an attention head.
num_heads: The number of attention heads to use. Defaults to 1.
num_blocks: Number of times to compute attention per time step. Defaults
to 1.
forget_bias: Bias to use for the forget gate, assuming we are using
some form of gating. Defaults to 1.
input_bias: Bias to use for the input gate, assuming we are using
some form of gating. Defaults to 0.
gate_style: Whether to use per-element gating ('unit'),
per-memory slot gating ('memory'), or no gating at all (None).
Defaults to `unit`.
attention_mlp_layers: Number of layers to use in the post-attention
MLP. Defaults to 2.
key_size: Size of vector to use for key & query vectors in the attention
computation. Defaults to None, in which case we use `head_size`.
name: Name of the module.
Raises:
ValueError: gate_style not one of [None, 'memory', 'unit'].
ValueError: num_blocks is < 1.
ValueError: attention_mlp_layers is < 1.
"""
self._mem_slots = mem_slots
self._head_size = head_size
self._num_heads = num_heads
self._mem_size = self._head_size * self._num_heads
self._name = name
if num_blocks < 1:
raise ValueError(
'num_blocks must be >= 1. Got: {}.'.format(num_blocks))
self._num_blocks = num_blocks
self._forget_bias = forget_bias
self._input_bias = input_bias
if gate_style not in ['unit', 'memory', None]:
raise ValueError(
'gate_style must be one of [\'unit\', \'memory\', None]. Got: '
'{}.'.format(gate_style))
self._gate_style = gate_style
if attention_mlp_layers < 1:
raise ValueError(
'attention_mlp_layers must be >= 1. Got: {}.'.format(
attention_mlp_layers))
self._attention_mlp_layers = attention_mlp_layers
self._key_size = key_size if key_size else self._head_size
self._template = tf.make_template(
self._name, self._build) # wrapper for variable sharing
with tf.variable_scope("sample_z"):
scale = tf.sqrt(tf.exp(log_var))
dist = tf.distributions.Normal(loc=loc, scale=scale)
z = dist.sample()
return z
def vae_model(x, z_dim):
loc, log_var = inference_network(x)
z = sample_z(loc, log_var)
x_hat = generative_network(z)
return loc, log_var, z, x_hat
model_opt = {"z_dim": FLAGS.z_dim}
model = tf.make_template('vae', vae_model)
xs = [
tf.placeholder(tf.float32, shape=(None, 128, 128, 3))
for i in range(FLAGS.nr_gpu)
]
ms = [
tf.placeholder_with_default(np.ones((FLAGS.batch_size, 128, 128),
dtype=np.float32),
shape=(None, 128, 128))
for i in range(FLAGS.nr_gpu)
]
mxs = tf.multiply(xs, tf.stack([ms for k in range(3)], axis=-1))
zs = [
tf.placeholder(tf.float32, shape=(None, 100)) for i in range(FLAGS.nr_gpu)
]
def build_model(self):
seed = 0
np.random.seed(seed)
tf.set_random_seed(seed)
image_dims = [self.input_height, self.input_width, self.c_dim]
self.inputs = tf.placeholder(tf.float32,
[self.batch_size] + image_dims,
name='real_images')
self.sample_inputs = tf.placeholder(tf.float32,
[self.sample_num] + image_dims,
name='sample_inputs')
self.image_size = np.prod(image_dims)
self.image_dims = image_dims
if self.dataset_name == "cifar":
inputs = tf.map_fn(
lambda img: tf.image.random_flip_left_right(img), self.inputs)
else:
inputs = self.inputs
sample_inputs = self.sample_inputs
self.z = tf.placeholder(tf.float32, [self.batch_size, self.z_dim],
name='z')
self.z_sum = histogram_summary("z", self.z)
#### f: Image Space to Latent space #########
#self.flow_model = tf.make_template('model',
# lambda x: nvp.model_spec(x, reuse=False, model_type=self.model_type, train=False,
# alpha=self.alpha, init_type=self.init_type, hidden_layers=self.hidden_layers,
# no_of_layers=self.no_of_layers, batch_norm_adaptive=1), unique_name_='model')
#self.flow_model = tf.make_template('model',
# lambda x: nvp.model_spec(x, reuse=False, model_type=self.model_type, train=False,
# alpha=self.alpha, init_type=self.init_type,
# hidden_layers=self.hidden_layers,
# no_of_layers=self.no_of_layers, batch_norm_adaptive=1),
# unique_name_='model')
#self.flow_model = tf.make_template('model', lambda x: nvp.model_spec(x, reuse=True, model_type=self.model_type, train=False,
# alpha=self.alpha, init_type=self.init_type,
# hidden_layers=self.hidden_layers,
# no_of_layers=self.no_of_layers, batch_norm_adaptive=1),
# unique_name_='model')
self.flow_model = tf.make_template(
'model',
lambda x: nvp.model_spec(x,
reuse=False,
model_type=self.model_type,
train=False,
alpha=self.alpha,
init_type=self.init_type,
hidden_layers=self.hidden_layers,
no_of_layers=self.no_of_layers,
batch_norm_adaptive=1),
unique_name_='model')
#### f: Image Space to Latent space for training #########
self.trainable_flow_model = tf.make_template(
'model',
lambda x: nvp.model_spec(x,
reuse=True,
model_type=self.model_type,
train=True,
alpha=self.alpha,
init_type=self.init_type,
hidden_layers=self.hidden_layers,
no_of_layers=self.no_of_layers,
batch_norm_adaptive=1),
unique_name_='model')
# ##### f^-1: Latent to image (trainable)#######
self.flow_inv_model = tf.make_template(
'model',
lambda x: nvp.inv_model_spec(x,
reuse=True,
model_type=self.model_type,
train=True,
alpha=self.alpha),
unique_name_='model')
# ##### f^-1: Latent to image (not-trainable just for sampling)#######
self.sampler_function = tf.make_template(
'model',
lambda x: nvp.inv_model_spec(x,
reuse=True,
model_type=self.model_type,
alpha=self.alpha,
train=False),
unique_name_='model')
self.generator_train_batch = self.flow_inv_model
############### SET SIZE FOR TEST BATCH DEPENDING ON WHETHER WE USE Linear or Conv arch##########
if self.model_type == "nice":
self.log_like_batch = tf.placeholder(\
tf.float32, [self.batch_size, self.image_size], name='log_like_batch')
elif self.model_type == "real_nvp":
self.log_like_batch = tf.placeholder(\
tf.float32, [self.batch_size] + self.image_dims, name='log_like_batch')
###############################################
gen_para, jac = self.flow_model(self.log_like_batch)
if self.dataset_name == "mnist":
self.log_likelihood = nvp_op.log_likelihood(
gen_para, jac, self.prior) / (self.batch_size)
else:
# to calculate values in bits per dim we need to
# multiply the density by the width of the
# discrete probability area, which is 1/256.0, per dimension.
# The calculation is performed in the log space.
self.log_likelihood = nvp_op.log_likelihood(
gen_para, jac, self.prior) / (self.batch_size)
self.log_likelihood = 8. + self.log_likelihood / (np.log(2) *
self.image_size)
self.G_before_postprocessing = self.generator_train_batch(self.z)
self.sampler_before_postprocessing = self.sampler_function(self.z)
if self.model_type == "real_nvp":
##For data dependent init (not completely implemented)
self.x_init = tf.placeholder(tf.float32,
shape=[self.batch_size] + image_dims)
# run once for data dependent initialization of parameters
self.trainable_flow_model(self.x_init)
inputs_tr_flow = inputs
if self.model_type == "nice":
split_val = int(self.image_size / 2)
self.permutation = np.arange(self.image_size)
tmp = self.permutation.copy()
self.permutation[:split_val] = tmp[::2]
self.permutation[split_val:] = tmp[1::2]
self.for_perm = np.identity(self.image_size)
self.for_perm = tf.constant(self.for_perm[:, self.permutation],
tf.float32)
self.rev_perm = np.identity(self.image_size)
self.rev_perm = tf.constant(
self.rev_perm[:, np.argsort(self.permutation)], tf.float32)
self.G_before_postprocessing \
= tf.matmul(self.G_before_postprocessing,self.rev_perm)
self.sampler_before_postprocessing \
= tf.clip_by_value(tf.matmul(self.sampler_before_postprocessing, self.rev_perm) , 0., 1.)
inputs_tr_flow = tf.matmul(
tf.reshape(inputs, [self.batch_size, self.image_size]),
self.for_perm)
train_gen_para, train_jac = self.trainable_flow_model(inputs_tr_flow)
self.train_log_likelihood = nvp_op.log_likelihood(
train_gen_para, train_jac, self.prior) / self.batch_size
self.sampler = tf.reshape(self.sampler_before_postprocessing,
[self.batch_size] + image_dims)
self.G = tf.reshape(self.G_before_postprocessing,
[self.batch_size] + image_dims)
inputs = inputs * 255.0
corruption_level = 1.0
inputs = inputs + corruption_level * tf.random_uniform(
[self.batch_size] + image_dims)
inputs = inputs / (255.0 + corruption_level)
self.D, self.D_logits = self.discriminator(inputs, reuse=False)
self.D_, self.D_logits_ = self.discriminator(self.G, reuse=True)
self.d_sum = histogram_summary("d", self.D)
self.d__sum = histogram_summary("d_", self.D_)
self.G_sum = image_summary("G", self.G)
def sigmoid_cross_entropy_with_logits(x, y):
try:
return tf.nn.sigmoid_cross_entropy_with_logits(logits=x,
labels=y)
except:
return tf.nn.sigmoid_cross_entropy_with_logits(logits=x,
targets=y)
### Vanilla gan loss
if self.f_div == 'ce':
self.d_loss_real = tf.reduce_mean(
sigmoid_cross_entropy_with_logits(self.D_logits,
tf.ones_like(self.D)))
self.d_loss_fake = tf.reduce_mean(
sigmoid_cross_entropy_with_logits(self.D_logits_,
tf.zeros_like(self.D_)))
self.g_loss = tf.reduce_mean(
sigmoid_cross_entropy_with_logits(self.D_logits_,
tf.ones_like(self.D_)))
else:
### other gan losses
if self.f_div == 'hellinger':
self.d_loss_real = tf.reduce_mean(tf.exp(-self.D_logits))
self.d_loss_fake = tf.reduce_mean(tf.exp(self.D_logits_) - 2.)
self.g_loss = tf.reduce_mean(tf.exp(-self.D_logits_))
elif self.f_div == 'rkl':
self.d_loss_real = tf.reduce_mean(tf.exp(self.D_logits))
self.d_loss_fake = tf.reduce_mean(-self.D_logits_ - 1.)
self.g_loss = -tf.reduce_mean(-self.D_logits_ - 1.)
elif self.f_div == 'kl':
self.d_loss_real = tf.reduce_mean(-self.D_logits)
self.d_loss_fake = tf.reduce_mean(tf.exp(self.D_logits_ - 1.))
self.g_loss = tf.reduce_mean(-self.D_logits_)
elif self.f_div == 'tv':
self.d_loss_real = tf.reduce_mean(-0.5 *
tf.tanh(self.D_logits))
self.d_loss_fake = tf.reduce_mean(0.5 *
tf.tanh(self.D_logits_))
self.g_loss = tf.reduce_mean(-0.5 * tf.tanh(self.D_logits_))
elif self.f_div == 'lsgan':
self.d_loss_real = 0.5 * tf.reduce_mean((self.D_logits - 1)**2)
self.d_loss_fake = 0.5 * tf.reduce_mean(self.D_logits_**2)
self.g_loss = 0.5 * tf.reduce_mean((self.D_logits_ - 1)**2)
elif self.f_div == "wgan":
self.g_loss = -tf.reduce_mean(self.D_logits_)
self.d_loss_real = -tf.reduce_mean(self.D_logits)
self.d_loss_fake = tf.reduce_mean(self.D_logits_)
alpha = tf.random_uniform(shape=[1, self.batch_size],
minval=0.,
maxval=1.)
fake_data = self.G
real_data = inputs
differences = fake_data - real_data
interpolates = real_data + \
tf.transpose(alpha*tf.transpose(differences, perm=[1,2,3,0]), [3,0,1,2])
_, d_inter = self.discriminator(interpolates, reuse=True)
gradients = tf.gradients(d_inter, [interpolates])[0]
slopes = tf.sqrt(
tf.reduce_sum(tf.square(gradients), reduction_indices=[1]))
self.gradient_penalty = tf.reduce_mean((slopes - 1.)**2)
else:
print("ERROR: Unrecognized f-divergence...exiting")
exit(-1)
self.d_loss_real_sum = scalar_summary("d_loss_real", self.d_loss_real)
self.d_loss_fake_sum = scalar_summary("d_loss_fake", self.d_loss_fake)
if self.f_div == "wgan":
self.d_loss = self.d_loss_real + self.d_loss_fake + self.reg * self.gradient_penalty
else:
self.d_loss = self.d_loss_real + self.d_loss_fake
self.g_loss_sum = scalar_summary("g_loss", self.g_loss)
self.d_loss_sum = scalar_summary("d_loss", self.d_loss)
t_vars = tf.trainable_variables()
self.d_vars = [var for var in t_vars if '/d_' in var.name]
self.g_vars = [var for var in t_vars if '/g_' in var.name]
print("gen_vars:")
for var in self.g_vars:
print(var.name)
print("disc_vars:")
for var in self.d_vars:
print(var.name)
self.saver = tf.train.Saver(max_to_keep=0)
def create_graph(training_data_set, testing_data_set, validation_data_set,
class_range, batch_size, device_id, num_epochs,
algorithm_params, model, augmentation_info,
create_separate_validation_branch):
deep_nn_template = tf.make_template('nn_core',
model.create_tensor_graph,
class_count=class_range.stop)
####################################################################################
training_input_iter = training_nn_iterator(training_data_set,
augmentation_info, batch_size,
num_epochs, device_id)
images, labels = training_input_iter.get_next()
training_y_conv, cross_entropy, learning_rate, train_step = optimize_nn(
deep_nn_template,
images,
labels,
device_id=device_id,
name_prefix='training',
algorithm_params=algorithm_params,
loss_func=model.get_loss_func)
train_nn_params = NNParams(input_iterator=training_input_iter,
data_with_labels=None,
metrics=None,
predict_tensor=None)
####################################################################################
testing_input_iter = simple_nn_iterator(testing_data_set, batch_size)
testing_images, testing_labels = testing_input_iter.get_next()
model_input_params = ModelInputParams(x=testing_images,
y=None,
device_id=device_id,
is_training=False)
testing_tensor_outputs = deep_nn_template(
model_input_params, algorithm_params=algorithm_params)
test_metric_ops_holder = create_metrics(testing_labels,
testing_tensor_outputs.y_conv,
class_range, 'testing')
testing_nn_params = NNParams(input_iterator=testing_input_iter,
data_with_labels=None,
metrics=test_metric_ops_holder,
predict_tensor=None)
####################################################################################
validation_nn_params = NNParams(input_iterator=testing_input_iter,
data_with_labels=None,
metrics=test_metric_ops_holder,
predict_tensor=None)
if create_separate_validation_branch:
validation_input_iter = simple_nn_iterator(validation_data_set,
batch_size)
validation_images, validation_labels = validation_input_iter.get_next()
validation_model_input_params = ModelInputParams(x=validation_images,
y=None,
device_id=device_id,
is_training=False)
validation_tensor_outputs = deep_nn_template(
validation_model_input_params, algorithm_params=algorithm_params)
validation_metric_ops_holder = create_metrics(
validation_labels, validation_tensor_outputs.y_conv, class_range,
'validation')
validation_nn_params = NNParams(input_iterator=validation_input_iter,
data_with_labels=None,
metrics=validation_metric_ops_holder,
predict_tensor=None)
####################################################################################
return cross_entropy, learning_rate, testing_nn_params, train_nn_params, validation_nn_params, train_step
def __init__(self,
policy,
env,
nsteps,
ent_coef=0.01,
vf_coef=0.5,
max_grad_norm=0.5,
lr=7e-4,
alpha=0.99,
epsilon=1e-5,
total_timesteps=int(80e6),
lrschedule='linear',
vi_coef=1,
iterative=False):
sess = tf_util.get_session()
nenvs = env.num_envs
nbatch = nenvs * nsteps
with tf.variable_scope('a2c_model', reuse=tf.AUTO_REUSE):
op_rollout = tf.make_template(
'rollout', rollout_step,
n_actions=env.action_space.n) if iterative else None
# step_model is used for sampling
step_model = policy(nenvs, 1, sess, op_rollout)
# train_model is used to train our network
train_model = policy(nbatch, nsteps, sess, op_rollout)
A = tf.placeholder(train_model.action.dtype, train_model.action.shape)
ADV = tf.placeholder(tf.float32, [nbatch])
R_sum = tf.placeholder(tf.float32, [nbatch]) # accumulated reward
R_onestep = tf.placeholder(tf.float32, [nbatch]) # one step reward
DONE = tf.placeholder(tf.float32, [nbatch])
LR = tf.placeholder(tf.float32, [])
# Calculate the loss
# Total loss = Policy gradient loss - entropy * entropy coefficient + Value coefficient * value loss
# Policy loss
neglogpac = train_model.pd.neglogp(A)
# L = A(s,a) * -logpi(a|s)
pg_loss = tf.reduce_mean(ADV * neglogpac)
# Entropy is used to improve exploration by limiting the premature convergence to suboptimal policy.
entropy = tf.reduce_mean(train_model.pd.entropy())
# Value loss
vf_loss = losses.mean_squared_error(tf.squeeze(train_model.vf), R_sum)
# value iteration loss
vi_loss = tf.constant(0.0)
if train_model.iterative:
vi_loss = training_loss(op_rollout, train_model.vi_state, A, R_sum,
R_onestep, nenvs, nsteps)
# with tf.variable_scope('a2c_model/rollout/value', reuse=True):
# value_weight = tf.math.reduce_sum(tf.get_variable('kernel'))
# else:
# with tf.variable_scope('a2c_model/vf', reuse=True):
# value_weight = tf.math.reduce_sum(tf.get_variable('w'))
# loss sum
if train_model.iterative:
loss = pg_loss - entropy * ent_coef + vf_loss * vf_coef + vi_loss * vi_coef
else:
loss = pg_loss - entropy * ent_coef + vf_loss * vf_coef
# Update parameters using loss
# 1. Get the model parameters
params = find_trainable_variables("a2c_model")
# 2. Calculate the gradients
grads = tf.gradients(loss, params)
if max_grad_norm is not None:
# Clip the gradients (normalize)
grads, grad_norm = tf.clip_by_global_norm(grads, max_grad_norm)
grads = list(zip(grads, params))
# zip aggregate each gradient with parameters associated
# For instance zip(ABCD, xyza) => Ax, By, Cz, Da
# 3. Make op for one policy and value update step of A2C
trainer = tf.train.RMSPropOptimizer(learning_rate=LR,
decay=alpha,
epsilon=epsilon)
_train = trainer.apply_gradients(grads)
lr = Scheduler(v=lr, nvalues=total_timesteps, schedule=lrschedule)
def train(obs, states, sum_rewards, masks, actions, values, dones,
onestep_rewards):
# Here we calculate advantage A(s,a) = R + yV(s') - V(s)
# rewards = R + yV(s')
advs = sum_rewards - values
for step in range(len(obs)):
cur_lr = lr.value()
td_map = {
train_model.X: obs,
A: actions,
ADV: advs,
R_sum: sum_rewards,
DONE: dones,
R_onestep: onestep_rewards,
LR: cur_lr
}
if states is not None:
td_map[train_model.S] = states
td_map[train_model.M] = masks
policy_loss, value_loss, policy_entropy, value_iteration_loss, _ = sess.run(
[pg_loss, vf_loss, entropy, vi_loss, _train], td_map)
return policy_loss, value_loss, policy_entropy, value_iteration_loss
self.train = train
self.train_model = train_model
self.step_model = step_model
self.step = step_model.step
self.value = step_model.value
self.initial_state = step_model.initial_state
self.save = functools.partial(tf_util.save_variables, sess=sess)
self.load = functools.partial(tf_util.load_variables, sess=sess)
tf.global_variables_initializer().run(session=sess)
def __init__(self,
train_batch_size=4096,
test_chain_batch_size=4096,
bijector="iaf",
log_dir="/tmp/neutra",
base_learning_rate=1e-3,
q_base_scale=1.,
learning_rate_schedule=[[6000, 1e-1]]):
target, target_spec = GetTargetSpec()
self.target = target
self.target_spec = target_spec
with gin.config_scope("train"):
train_target, train_target_spec = GetTargetSpec()
self.train_target = train_target
self.train_target_spec = train_target_spec
if bijector == "rnvp":
bijector_fn = tf.make_template("bijector",
MakeRNVPBijectorFn,
num_dims=self.target_spec.num_dims)
elif bijector == "iaf":
bijector_fn = tf.make_template("bijector",
MakeIAFBijectorFn,
num_dims=self.target_spec.num_dims)
elif bijector == "affine":
bijector_fn = tf.make_template("bijector",
MakeAffineBijectorFn,
num_dims=self.target_spec.num_dims)
else:
bijector_fn = lambda *args, **kwargs: tfb.Identity()
self.train_bijector = bijector_fn(train=True)
self.bijector = bijector_fn(train=False)
if train_target_spec.bijector is not None:
print("Using train target bijector")
self.train_bijector = tfb.Chain(
[train_target_spec.bijector, self.train_bijector])
if target_spec.bijector is not None:
print("Using target bijector")
self.bijector = tfb.Chain([target_spec.bijector, self.bijector])
q_base = tfd.Independent(
tfd.Normal(loc=tf.zeros(self.target_spec.num_dims),
scale=q_base_scale *
tf.ones(self.target_spec.num_dims)), 1)
self.q_x_train = tfd.TransformedDistribution(q_base,
self.train_bijector)
self.q_x = tfd.TransformedDistribution(q_base, self.bijector)
# Params
self.train_batch_size = int(train_batch_size)
self.test_chain_batch_size = tf.placeholder_with_default(
test_chain_batch_size, [], "test_chain_batch_size")
self.test_batch_size = tf.placeholder_with_default(
16384 * 8, [], "test_batch_size")
self.test_num_steps = tf.placeholder_with_default(
1000, [], "test_num_steps")
self.test_num_leapfrog_steps = tf.placeholder_with_default(
tf.to_int64(2), [], "test_num_leapfrog_steps")
self.test_step_size = tf.placeholder_with_default(
0.1, [], "test_step_size")
# Test
self.neutra_outputs = MakeNeuTra(
target=self.target,
q=self.q_x,
batch_size=self.test_chain_batch_size,
num_steps=self.test_num_steps,
num_leapfrog_steps=self.test_num_leapfrog_steps,
step_size=self.test_step_size,
)
self.z_chain = tf.reshape(
self.bijector.inverse(
tf.reshape(self.neutra_outputs.x_chain,
[-1, self.target_spec.num_dims])),
tf.shape(self.neutra_outputs.x_chain))
self.target_samples = self.target.sample(self.test_batch_size)
self.target_z = self.bijector.inverse(self.target_samples)
self.q_samples = self.q_x.sample(self.test_batch_size)
self.target_cov = utils.Covariance(self.target_samples)
self.target_eigvals, self.target_eigvecs = tf.linalg.eigh(
self.target_cov)
self.cached_target_eigvals = tf.get_local_variable(
"cached_target_eigvals",
self.target_eigvals.shape,
initializer=tf.zeros_initializer())
self.cached_target_eigvecs = tf.get_local_variable(
"cached_target_eigvecs",
self.target_eigvecs.shape,
initializer=tf.zeros_initializer())
self.cached_target_stats_update_op = [
self.cached_target_eigvals.assign(self.target_eigvals),
self.cached_target_eigvecs.assign(self.target_eigvecs),
tf.print("Assigning target stats")
]
def variance(x):
x -= tf.reduce_mean(x, 0, keep_dims=True)
x = tf.square(x)
return x
def rotated_variance(x):
x2 = tf.reshape(x, [-1, self.target_spec.num_dims])
x2 -= tf.reduce_mean(x2, 0, keep_dims=True)
x2 = tf.matmul(x2, self.cached_target_eigvecs)
x2 = tf.square(x2)
return tf.reshape(x2, tf.shape(x))
functions = [
("mean", tf.identity),
# ("var", variance),
("square", tf.square),
# ("rot_square", rot_square),
# ("rot_var", rotated_variance),
]
self.cached_target_mean = {}
self.cached_target_mean_update_op = [
tf.print("Assigning target means.")
]
self.neutra_stats = {}
self.q_stats = {}
for name, f in functions:
target_mean = tf.reduce_mean(f(self.target_samples), 0)
cached_target_mean = tf.get_local_variable(name + "_cached_mean",
target_mean.shape)
if self.target_spec.stats is not None:
self.cached_target_mean_update_op.append(
cached_target_mean.assign(self.target_spec.stats[name]))
else:
self.cached_target_mean_update_op.append(
cached_target_mean.assign(target_mean))
self.cached_target_mean[name] = cached_target_mean
self.q_stats[name] = ComputeQStats(f(self.q_samples),
cached_target_mean)
self.neutra_stats[name] = ComputeChainStats(
f(self.neutra_outputs.x_chain), cached_target_mean,
self.test_num_leapfrog_steps)
# Training
self.train_q_samples = self.q_x_train.sample(self.train_batch_size)
self.train_log_q_x = self.q_x_train.log_prob(self.train_q_samples)
self.kl_q_p = tf.reduce_mean(
self.train_log_q_x - self.target.log_prob(self.train_q_samples))
loss = self.kl_q_p
reg_losses = tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES)
if reg_losses:
tf.logging.info("Regularizing.")
loss += tf.add_n(reg_losses)
self.loss = tf.check_numerics(loss, "Loss has NaNs")
self.global_step = tf.train.get_or_create_global_step()
steps, factors = zip(*learning_rate_schedule)
learning_rate = base_learning_rate * tf.train.piecewise_constant(
self.global_step, steps, [1.0] + list(factors))
opt = tf.train.AdamOptimizer(learning_rate=learning_rate)
self.train_op = opt.minimize(self.loss, global_step=self.global_step)
tf.contrib.summary.scalar("kl_q_p", self.kl_q_p)
tf.contrib.summary.scalar("loss", self.loss)
self.init = [
tf.global_variables_initializer(),
tf.local_variables_initializer(),
tf.print("Initializing variables")
]
self.saver = tf.train.Saver()
self.log_dir = log_dir
def _build_networks(self):
"""Builds the IQN computations needed for acting and training.
These are:
self.online_convnet: For computing the current state's quantile values.
self.target_convnet: For computing the next state's target quantile
values.
self._net_outputs: The actual quantile values.
self._q_argmax: The action maximizing the current state's Q-values.
self._replay_net_outputs: The replayed states' quantile values.
self._replay_next_target_net_outputs: The replayed next states' target
quantile values.
"""
# Calling online_convnet will generate a new graph as defined in
# self._get_network_template using whatever input is passed, but will always
# share the same weights.
self.online_convnet = tf.make_template('Online',
self._network_template)
self.target_convnet = tf.make_template('Target',
self._network_template)
# Compute the Q-values which are used for action selection in the current
# state.
self._net_outputs = self.online_convnet(self.state_ph,
self.num_quantile_samples)
# Shape of self._net_outputs.quantile_values:
# num_quantile_samples x num_actions.
# e.g. if num_actions is 2, it might look something like this:
# Vals for Quantile .2 Vals for Quantile .4 Vals for Quantile .6
# [[0.1, 0.5], [0.15, -0.3], [0.15, -0.2]]
# Q-values = [(0.1 + 0.15 + 0.15)/3, (0.5 + 0.15 + -0.2)/3].
self._q_values = tf.reduce_mean(self._net_outputs.quantile_values,
axis=0)
self._q_argmax = tf.argmax(self._q_values, axis=0)
self._replay_net_outputs = self.online_convnet(self._replay.states,
self.num_tau_samples)
# Shape: (num_tau_samples x batch_size) x num_actions.
self._replay_net_quantile_values = self._replay_net_outputs.quantile_values
self._replay_net_quantiles = self._replay_net_outputs.quantiles
# Do the same for next states in the replay buffer.
self._replay_net_target_outputs = self.target_convnet(
self._replay.next_states, self.num_tau_prime_samples)
# Shape: (num_tau_prime_samples x batch_size) x num_actions.
vals = self._replay_net_target_outputs.quantile_values
self._replay_net_target_quantile_values = vals
# Compute Q-values which are used for action selection for the next states
# in the replay buffer. Compute the argmax over the Q-values.
if self.double_dqn:
outputs_action = self.online_convnet(self._replay.next_states,
self.num_quantile_samples)
else:
outputs_action = self.target_convnet(self._replay.next_states,
self.num_quantile_samples)
# Shape: (num_quantile_samples x batch_size) x num_actions.
target_quantile_values_action = outputs_action.quantile_values
# Shape: num_quantile_samples x batch_size x num_actions.
target_quantile_values_action = tf.reshape(
target_quantile_values_action, [
self.num_quantile_samples, self._replay.batch_size,
self.num_actions
])
# Shape: batch_size x num_actions.
self._replay_net_target_q_values = tf.squeeze(
tf.reduce_mean(target_quantile_values_action, axis=0))
self._replay_next_qt_argmax = tf.argmax(
self._replay_net_target_q_values, axis=1)
def get_generator(model_name, scope='generator', **kwargs):
model_func = generator_dict[model_name]
return tf.make_template(scope, model_func, **kwargs)
configs_dir = __file__.split('/')[-2]
config = importlib.import_module('%s.%s' % (configs_dir, args.config_name))
save_dir = utils.find_model_metadata('metadata/', args.config_name)
experiment_id = os.path.dirname(save_dir)
if not os.path.isdir(save_dir + '/samples'):
os.makedirs(save_dir + '/samples')
samples_dir = save_dir + '/samples'
print('exp_id', experiment_id)
print('n_context', config.n_context)
print('seq_len', config.seq_len)
# create the model
model = tf.make_template('model', config.build_model, sampling_mode=True)
all_params = tf.trainable_variables()
x_in = tf.placeholder(tf.float32, shape=(1, ) + config.obs_shape)
y_label = tf.placeholder(tf.float32, shape=(1, ) + config.label_shape)
samples = model(x_in, y_label)
saver = tf.train.Saver()
with tf.Session() as sess:
begin = time.time()
ckpt_file = save_dir + 'params.ckpt'
print('restoring parameters from', ckpt_file)
saver.restore(sess, tf.train.latest_checkpoint(save_dir))
prior_samples = []
def run():
(x_train, y_train), (x_test, y_test) = observations.mnist('~/data/MNIST')
dataset_size = x_train.shape[0]
inputs = tf.placeholder(tf.float32, [None, 28 * 28], 'inputs')
labels = tf.placeholder(tf.int32, [None], 'labels')
data_indices = tf.placeholder(tf.int32, [None], 'data_indices')
def network(context: modular.ModularContext):
modules = modular.create_dense_modules(inputs, module_count=10, units=32)
hidden = modular.modular_layer(inputs, modules, parallel_count=1, context=context)
hidden = tf.nn.relu(hidden)
modules = modular.create_dense_modules(hidden, module_count=8, units=10)
logits = modular.modular_layer(hidden, modules, parallel_count=1, context=context)
target = modular.modularize_target(labels, context)
loglikelihood = tf.distributions.Categorical(logits).log_prob(target)
predicted = tf.argmax(logits, axis=-1, output_type=tf.int32)
accuracy = tf.reduce_mean(tf.cast(tf.equal(predicted, target), tf.float32))
selection_entropy = context.selection_entropy()
batch_selection_entropy = context.batch_selection_entropy()
return loglikelihood, logits, accuracy, selection_entropy, batch_selection_entropy
template = tf.make_template('network', network)
optimizer = tf.train.AdamOptimizer()
e_step, m_step, eval = modular.modularize(template, optimizer, dataset_size,
data_indices, sample_size=10)
ll, logits, accuracy, s_entropy, bs_entropy = eval
tf.summary.scalar('loglikelihood', tf.reduce_mean(ll))
tf.summary.scalar('accuracy', accuracy)
tf.summary.scalar('entropy/exp_selection', tf.exp(s_entropy))
tf.summary.scalar('entropy/exp_batch_selection', tf.exp(bs_entropy))
with tf.Session() as sess:
time = '{:%Y-%m-%d_%H:%M:%S}'.format(datetime.datetime.now())
writer = tf.summary.FileWriter('logs/train_{}'.format(time))
test_writer = tf.summary.FileWriter('logs/test_{}'.format(time))
general_summaries = tf.summary.merge_all()
m_step_summaries = tf.summary.merge([create_m_step_summaries(), general_summaries])
sess.run(tf.global_variables_initializer())
batches = generator([x_train, y_train, np.arange(dataset_size)], 32)
for i, (batch_x, batch_y, indices) in enumerate(batches):
feed_dict = {
inputs: batch_x,
labels: batch_y,
data_indices: indices
}
step = e_step if i % 10 == 0 else m_step
summaries = m_step_summaries if step == m_step else general_summaries
_, summary_data = sess.run([step, summaries], feed_dict)
writer.add_summary(summary_data, global_step=i)
if i % 100 == 0:
test_feed_dict = {inputs: x_test, labels: y_test}
summary_data = sess.run(general_summaries, test_feed_dict)
test_writer.add_summary(summary_data, global_step=i)
writer.close()
test_writer.close()
def build_iteration(self,
iteration_number,
ensemble_candidates,
subnetwork_builders,
features,
mode,
labels=None,
previous_ensemble_summary=None,
previous_ensemble_spec=None,
rebuilding=False):
"""Builds and returns AdaNet iteration t.
This method uses the generated the candidate subnetworks given the ensemble
at iteration t-1 and creates graph operations to train them. The returned
`_Iteration` tracks the training of all candidates to know when the
iteration is over, and tracks the best candidate's predictions and loss, as
defined by lowest complexity-regularized loss on the train set.
Args:
iteration_number: Integer iteration number.
ensemble_candidates: Iterable of `adanet.ensemble.Candidate` instances.
subnetwork_builders: A list of `Builders` for adding ` Subnetworks` to the
graph. Each subnetwork is then wrapped in a `_Candidate` to train.
features: Dictionary of `Tensor` objects keyed by feature name.
mode: Defines whether this is training, evaluation or prediction. See
`ModeKeys`.
labels: `Tensor` of labels. Can be `None`.
previous_ensemble_summary: The `adanet.Summary` for the previous ensemble.
previous_ensemble_spec: Optional `_EnsembleSpec` for iteration t-1.
rebuilding: Boolean whether the iteration is being rebuilt only to restore
the previous best subnetworks and ensembles.
Returns:
An _Iteration instance.
Raises:
ValueError: If subnetwork_builders is empty.
ValueError: If two subnetworks share the same name.
ValueError: If two ensembles share the same name.
"""
tf.logging.info("%s iteration %s",
"Rebuilding" if rebuilding else "Building",
iteration_number)
if not subnetwork_builders:
raise ValueError("Each iteration must have at least one Builder.")
# TODO: Consider moving builder mode logic to ensemble_builder.py.
builder_mode = mode
if rebuilding:
# Build the subnetworks and ensembles in EVAL mode by default. This way
# their outputs aren't affected by dropout etc.
builder_mode = tf.estimator.ModeKeys.EVAL
if mode == tf.estimator.ModeKeys.PREDICT:
builder_mode = mode
# Only replicate in training mode when the user requests it.
if self._replicate_ensemble_in_training and (
mode == tf.estimator.ModeKeys.TRAIN):
builder_mode = mode
features, labels = self._check_numerics(features, labels)
training = mode == tf.estimator.ModeKeys.TRAIN
skip_summaries = mode == tf.estimator.ModeKeys.PREDICT
with tf.variable_scope("iteration_{}".format(iteration_number)):
# Iteration step to use instead of global step.
iteration_step = tf.get_variable(
"step",
shape=[],
initializer=tf.zeros_initializer(),
trainable=False,
dtype=tf.int64)
# Convert to tensor so that users cannot mutate it.
iteration_step_tensor = tf.convert_to_tensor(iteration_step)
seen_builder_names = {}
candidates = []
summaries = []
subnetwork_reports = {}
previous_ensemble = None
if previous_ensemble_spec:
previous_ensemble = previous_ensemble_spec.ensemble
# Include previous best subnetwork as a candidate so that its
# predictions are returned until a new candidate outperforms.
seen_builder_names = {previous_ensemble_spec.name: True}
previous_best_candidate = self._candidate_builder.build_candidate(
ensemble_spec=previous_ensemble_spec,
training=training,
iteration_step=iteration_step_tensor,
summary=previous_ensemble_summary,
is_previous_best=True)
candidates.append(previous_best_candidate)
summaries.append(previous_ensemble_summary)
# Generate subnetwork reports.
if mode == tf.estimator.ModeKeys.EVAL:
metrics = call_eval_metrics(
previous_ensemble_spec.eval_metrics)
subnetwork_report = subnetwork.Report(
hparams={},
attributes={},
metrics=metrics,
)
subnetwork_report.metrics["adanet_loss"] = tf.metrics.mean(
previous_ensemble_spec.adanet_loss)
subnetwork_reports["previous_ensemble"] = subnetwork_report
for subnetwork_builder in subnetwork_builders:
if subnetwork_builder.name in seen_builder_names:
raise ValueError(
"Two subnetworks have the same name '{}'".format(
subnetwork_builder.name))
seen_builder_names[subnetwork_builder.name] = True
subnetwork_specs = []
num_subnetworks = len(subnetwork_builders)
for i, subnetwork_builder in enumerate(subnetwork_builders):
if not self._placement_strategy.should_build_subnetwork(
num_subnetworks, i) and not rebuilding:
continue
subnetwork_name = "t{}_{}".format(iteration_number,
subnetwork_builder.name)
subnetwork_summary = self._summary_maker(
namespace="subnetwork",
scope=subnetwork_name,
skip_summary=skip_summaries or rebuilding)
summaries.append(subnetwork_summary)
tf.logging.info("%s subnetwork '%s'",
"Rebuilding" if rebuilding else "Building",
subnetwork_builder.name)
subnetwork_spec = self._subnetwork_manager.build_subnetwork_spec(
name=subnetwork_name,
subnetwork_builder=subnetwork_builder,
iteration_step=iteration_step_tensor,
summary=subnetwork_summary,
features=features,
mode=builder_mode,
labels=labels,
previous_ensemble=previous_ensemble)
subnetwork_specs.append(subnetwork_spec)
if not self._placement_strategy.should_build_ensemble(
num_subnetworks) and not rebuilding:
# Workers that don't build ensembles need a dummy candidate in order
# to train the subnetwork.
# Because only ensembles can be considered candidates, we need to
# convert the subnetwork into a dummy ensemble and subsequently a
# dummy candidate. However, this dummy candidate is never considered a
# true candidate during candidate evaluation and selection.
# TODO: Eliminate need for candidates.
dummy_candidate = self._candidate_builder.build_candidate(
# pylint: disable=protected-access
ensemble_spec=ensemble_builder_lib._EnsembleSpec(
name=subnetwork_name,
ensemble=None,
architecture=None,
subnetwork_builders=subnetwork_builders,
predictions=subnetwork_spec.predictions,
loss=subnetwork_spec.loss,
adanet_loss=0.),
# pylint: enable=protected-access
training=training,
iteration_step=iteration_step_tensor,
summary=subnetwork_summary,
track_moving_average=False)
candidates.append(dummy_candidate)
# Generate subnetwork reports.
if mode != tf.estimator.ModeKeys.PREDICT:
subnetwork_report = subnetwork_builder.build_subnetwork_report(
)
if not subnetwork_report:
subnetwork_report = subnetwork.Report(hparams={},
attributes={},
metrics={})
metrics = call_eval_metrics(subnetwork_spec.eval_metrics)
for metric_name in sorted(metrics):
metric = metrics[metric_name]
subnetwork_report.metrics[metric_name] = metric
subnetwork_reports[
subnetwork_builder.name] = subnetwork_report
# Create (ensembler_candidate*ensembler) ensembles.
seen_ensemble_names = {}
for ensembler in self._ensemblers:
for ensemble_candidate in ensemble_candidates:
if not self._placement_strategy.should_build_ensemble(
num_subnetworks) and not rebuilding:
continue
ensemble_name = "t{}_{}_{}".format(iteration_number,
ensemble_candidate.name,
ensembler.name)
if ensemble_name in seen_ensemble_names:
raise ValueError(
"Two ensembles have the same name '{}'".format(
ensemble_name))
seen_ensemble_names[ensemble_name] = True
summary = self._summary_maker(namespace="ensemble",
scope=ensemble_name,
skip_summary=skip_summaries
or rebuilding)
summaries.append(summary)
ensemble_spec = self._ensemble_builder.build_ensemble_spec(
name=ensemble_name,
candidate=ensemble_candidate,
ensembler=ensembler,
subnetwork_specs=subnetwork_specs,
summary=summary,
features=features,
mode=builder_mode,
iteration_step=iteration_step_tensor,
iteration_number=iteration_number,
labels=labels,
previous_ensemble_spec=previous_ensemble_spec)
candidate = self._candidate_builder.build_candidate(
ensemble_spec=ensemble_spec,
training=training,
iteration_step=iteration_step_tensor,
summary=summary)
candidates.append(candidate)
# TODO: Move adanet_loss from subnetwork report to a new
# ensemble report, since the adanet_loss is associated with an
# ensemble, and only when using a ComplexityRegularizedEnsemblers.
# Keep adanet_loss in subnetwork report for backwards compatibility.
if len(ensemble_candidates) != len(subnetwork_builders):
continue
if len(ensemble_candidate.subnetwork_builders) > 1:
continue
if mode == tf.estimator.ModeKeys.PREDICT:
continue
builder_name = ensemble_candidate.subnetwork_builders[
0].name
subnetwork_reports[builder_name].metrics[
"adanet_loss"] = tf.metrics.mean(
ensemble_spec.adanet_loss)
# Dynamically select the outputs of best candidate.
best_candidate_index = self._best_candidate_index(candidates)
best_predictions = self._best_predictions(candidates,
best_candidate_index)
best_loss = self._best_loss(candidates, best_candidate_index, mode)
best_export_outputs = self._best_export_outputs(
candidates, best_candidate_index, mode, best_predictions)
# Hooks on TPU cannot depend on any graph `Tensors`. Instead the value of
# `is_over` is stored in a `Variable` that can later be retrieved from
# inside a training hook.
is_over_var_template = tf.make_template("is_over_var_template",
_is_over_var)
training_chief_hooks, training_hooks = (), ()
for subnetwork_spec in subnetwork_specs:
if not self._placement_strategy.should_train_subnetworks(
num_subnetworks) and not rebuilding:
continue
if not subnetwork_spec.train_op:
continue
training_chief_hooks += subnetwork_spec.train_op.chief_hooks or (
)
training_hooks += subnetwork_spec.train_op.hooks or ()
for candidate in candidates:
spec = candidate.ensemble_spec
if not spec.train_op:
continue
training_chief_hooks += spec.train_op.chief_hooks or ()
training_hooks += spec.train_op.hooks or ()
summary = self._summary_maker(namespace=None,
scope=None,
skip_summary=skip_summaries
or rebuilding)
summaries.append(summary)
with summary.current_scope():
summary.scalar("iteration/adanet/iteration", iteration_number)
summary.scalar("iteration_step/adanet/iteration_step",
iteration_step_tensor)
if best_loss is not None:
summary.scalar("loss", best_loss)
train_op = self._create_train_op(subnetwork_specs, candidates,
mode, iteration_step,
is_over_var_template,
num_subnetworks)
iteration_metrics = _IterationMetrics(candidates, subnetwork_specs)
if self._use_tpu:
estimator_spec = tf.contrib.tpu.TPUEstimatorSpec(
mode=mode,
predictions=best_predictions,
loss=best_loss,
train_op=train_op,
eval_metrics=iteration_metrics.best_eval_metrics_tuple(
best_candidate_index, mode),
export_outputs=best_export_outputs,
training_hooks=training_hooks)
else:
estimator_spec = tf.estimator.EstimatorSpec(
mode=mode,
predictions=best_predictions,
loss=best_loss,
train_op=train_op,
eval_metric_ops=iteration_metrics.best_eval_metric_ops(
best_candidate_index, mode),
export_outputs=best_export_outputs,
training_chief_hooks=training_chief_hooks,
training_hooks=training_hooks)
return _Iteration(number=iteration_number,
candidates=candidates,
subnetwork_specs=subnetwork_specs,
estimator_spec=estimator_spec,
best_candidate_index=best_candidate_index,
summaries=summaries,
is_over_fn=is_over_var_template,
subnetwork_reports=subnetwork_reports,
step=iteration_step_tensor)
def __init__(self, batch_env, step, is_training, should_log, config):
"""Create an instance of the PPO algorithm.
Args:
batch_env: In-graph batch environment.
step: Integer tensor holding the current training step.
is_training: Boolean tensor for whether the algorithm should train.
should_log: Boolean tensor for whether summaries should be returned.
config: Object containing the agent configuration as attributes.
"""
self._batch_env = batch_env
self._step = step
self._is_training = is_training
self._should_log = should_log
self._config = config
self._observ_filter = normalize.StreamingNormalize(self._batch_env.observ[0],
center=True,
scale=True,
clip=5,
name='normalize_observ')
self._reward_filter = normalize.StreamingNormalize(self._batch_env.reward[0],
center=False,
scale=True,
clip=10,
name='normalize_reward')
# Memory stores tuple of observ, action, mean, logstd, reward.
template = (self._batch_env.observ[0], self._batch_env.action[0], self._batch_env.action[0],
self._batch_env.action[0], self._batch_env.reward[0])
self._memory = memory.EpisodeMemory(template, config.update_every, config.max_length, 'memory')
self._memory_index = tf.Variable(0, False)
use_gpu = self._config.use_gpu and utility.available_gpus()
with tf.device('/gpu:0' if use_gpu else '/cpu:0'):
# Create network variables for later calls to reuse.
action_size = self._batch_env.action.shape[1].value
self._network = tf.make_template('network',
functools.partial(config.network, config, action_size))
output = self._network(
tf.zeros_like(self._batch_env.observ)[:, None], tf.ones(len(self._batch_env)))
with tf.variable_scope('ppo_temporary'):
self._episodes = memory.EpisodeMemory(template, len(batch_env), config.max_length,
'episodes')
if output.state is None:
self._last_state = None
else:
# Ensure the batch dimension is set.
tf.contrib.framework.nest.map_structure(
lambda x: x.set_shape([len(batch_env)] + x.shape.as_list()[1:]), output.state)
# pylint: disable=undefined-variable
self._last_state = tf.contrib.framework.nest.map_structure(
lambda x: tf.Variable(lambda: tf.zeros_like(x), False), output.state)
self._last_action = tf.Variable(tf.zeros_like(self._batch_env.action),
False,
name='last_action')
self._last_mean = tf.Variable(tf.zeros_like(self._batch_env.action),
False,
name='last_mean')
self._last_logstd = tf.Variable(tf.zeros_like(self._batch_env.action),
False,
name='last_logstd')
self._penalty = tf.Variable(self._config.kl_init_penalty, False, dtype=tf.float32)
self._optimizer = self._config.optimizer(self._config.learning_rate)
def main():
# define dataset object
train_set = dataset(train_path, batch_size)
# define input
train_input = tf.placeholder(
tf.float32,
shape=(image_size,image_size,batch_size)
)
train_label = tf.placeholder(
tf.float32,
shape=(image_size,image_size,batch_size)
)
# the model and loss
shared_model = tf.make_template('shared_model', model)
train_output, weights = shared_model(train_input)
loss = tf.reduce_sum(tf.nn.l2_loss(tf.substract(train_output, train_label)))
# normlization weight.
for w in weights:
loss += tf.nn.l2_loss(w)*1e-4
# record loss
tf.summary.scalar("loss", loss)
# training step and learning rate
global_step = tf.Variable(0, trainable=False)
learning_rate = tf.train.exponential_decay(
base_learning_rate,
global_step*batch_size,
train_set.instance_num*lr_step_size,
lr_decay,
staircase=True
)
tf.summary.scalar("learning_rate", learning_rate)
# Optimizer
optimizer = tf.train.AdamOptimizer(learning_rate)
opt = optimizer.minimize(loss, global_step=global_step)
saver = tf.train.Saver(weights, max_to_keep=0)
config = tf.ConfigProto()
# training
with tf.Session(config=config) as sess:
#TensorBoard open log with "tensorboard --logdir=logs"
if not os.path.exists('logs'):
os.mkdir('logs')
merged = tf.summary.merge_all()
file_writer = tf.summary.FileWriter('logs',sess.graph)
# var initializaion
tf.initialize_all_variables().run()
if model_path:
print "restore model..."
saver.restore(sess,model_path)
print "successfully restore previous model."
# train
for epoch in xrange(0,max_epoch):
for step in range(train_set.instance_num//batch_size):
data, label = train_set.next_batch()
feed_dict = {train_input : data, train_label : label}
_,l,output,lr,g_step = sess.run([opt, loss, train_output, learning_rate, global_step],feed_dict=feed_dict)
print "[epoch %2.4f] loss %.4f\t lr %.5f" % (epoch+(float(step)*batch_size/train_set.instance_num), np.sum(l)/batch_size, lr)
del data, label
saver.save(sess, "./checkpoints/VDSR_const_clip_0.01_epoch_%03d.ckpt" % epoch ,global_step=global_step)
def _build_networks(self):
"""Builds the IQN computations needed for acting and training.
These are:
self.online_convnet: For computing the current state's quantile values.
self.target_convnet: For computing the next state's target quantile
values.
self._net_outputs: The actual quantile values.
self._q_argmax: The action maximizing the current state's Q-values.
self._replay_net_outputs: The replayed states' quantile values.
self._replay_next_target_net_outputs: The replayed next states' target
quantile values.
"""
# Calling online_convnet will generate a new graph as defined in
# self._get_network_template using whatever input is passed, but will always
# share the same weights.
self.online_convnet = tf.make_template('Online', self._network_template)
self.target_convnet = tf.make_template('Target', self._network_template)
# Compute the Q-values which are used for action selection in the current
# state.
self._net_outputs = self.online_convnet(self.state_ph,
self.num_quantile_samples)
# Shape of self._net_outputs.quantile_values:
# num_quantile_samples x num_actions.
# e.g. if num_actions is 2, it might look something like this:
# Vals for Quantile .2 Vals for Quantile .4 Vals for Quantile .6
# [[0.1, 0.5], [0.15, -0.3], [0.15, -0.2]]
# Q-values = [(0.1 + 0.15 + 0.15)/3, (0.5 + 0.15 + -0.2)/3].
self._q_values = tf.reduce_mean(self._net_outputs.quantile_values, axis=0)
self._q_argmax = tf.argmax(self._q_values, axis=0)
self._replay_net_outputs = self.online_convnet(self._replay.states,
self.num_tau_samples)
# Shape: (num_tau_samples x batch_size) x num_actions.
self._replay_net_quantile_values = self._replay_net_outputs.quantile_values
self._replay_net_quantiles = self._replay_net_outputs.quantiles
# Do the same for next states in the replay buffer.
self._replay_net_target_outputs = self.target_convnet(
self._replay.next_states, self.num_tau_prime_samples)
# Shape: (num_tau_prime_samples x batch_size) x num_actions.
vals = self._replay_net_target_outputs.quantile_values
self._replay_net_target_quantile_values = vals
# Compute Q-values which are used for action selection for the next states
# in the replay buffer.
outputs_q = self.target_convnet(
self._replay.next_states, self.num_quantile_samples)
# Shape: (num_quantile_samples x batch_size) x num_actions.
target_quantile_values_q = outputs_q.quantile_values
# Shape: num_quantile_samples x batch_size x num_actions.
target_quantile_values_q = tf.reshape(target_quantile_values_q,
[self.num_quantile_samples,
self._replay.batch_size,
self.num_actions])
# Shape: batch_size x num_actions.
self._replay_net_target_q_values = tf.squeeze(tf.reduce_mean(
target_quantile_values_q, axis=0))
# Compute the argmax over target net Q-values using different quantile
# inputs.
outputs_action = self.target_convnet(self._replay.next_states,
self.num_quantile_samples)
# Shape: (num_quantile_samples x batch_size) x num_actions.
target_quantile_values_action = outputs_action.quantile_values
# Shape: num_quantile_samples x batch_size x num_actions.
target_quantile_values_action = tf.reshape(target_quantile_values_action,
[self.num_quantile_samples,
self._replay.batch_size,
self.num_actions])
# Shape: batch_size x num_actions.
target_q_values_action = tf.squeeze(tf.reduce_mean(
target_quantile_values_action, axis=0))
self._replay_next_qt_argmax = tf.argmax(target_q_values_action, axis=1)
for i in range(args.nr_gpu)
]
hs = [tf.one_hot(ys[i], num_labels) for i in range(args.nr_gpu)]
else:
h_init = None
h_sample = [None] * args.nr_gpu
hs = h_sample
# create the model
model_opt = {
'nr_resnet': args.nr_resnet,
'nr_filters': args.nr_filters,
'nr_logistic_mix': args.nr_logistic_mix,
'resnet_nonlinearity': args.resnet_nonlinearity
}
model = tf.make_template('model', model_spec)
# run once for data dependent initialization of parameters
gen_par = model(x_init,
h_init,
init=True,
dropout_p=args.dropout_p,
**model_opt)
# keep track of moving average
all_params = tf.trainable_variables()
ema = tf.train.ExponentialMovingAverage(decay=args.polyak_decay)
maintain_averages_op = tf.group(ema.apply(all_params))
# get loss gradients over multiple GPUs
grads = []
def __init__(self, arch, is_training=False):
self.arch = arch
self.is_training = is_training
self._decode = tf.make_template('Decoder', self._generator)
self._encode = tf.make_template('Encoder', self._encoder)
def __init__(self,
observation_spec,
action_spec,
actor_net=networks.actor_net,
critic_net=networks.critic_net,
td_errors_loss=tf.losses.huber_loss,
dqda_clipping=0.,
actions_regularizer=0.,
target_q_clipping=None,
residual_phi=0.0,
debug_summaries=False):
"""Constructs a DDPG agent.
Args:
observation_spec: A TensorSpec defining the observations.
action_spec: A BoundedTensorSpec defining the actions.
actor_net: A callable that creates the actor network. Must take the
following arguments: states, num_actions. Please see networks.actor_net
for an example.
critic_net: A callable that creates the critic network. Must take the
following arguments: states, actions. Please see networks.critic_net
for an example.
td_errors_loss: A callable defining the loss function for the critic
td error.
dqda_clipping: (float) clips the gradient dqda element-wise between
[-dqda_clipping, dqda_clipping]. Does not perform clipping if
dqda_clipping == 0.
actions_regularizer: A scalar, when positive penalizes the norm of the
actions. This can prevent saturation of actions for the actor_loss.
target_q_clipping: (tuple of floats) clips target q values within
(low, high) values when computing the critic loss.
residual_phi: (float) [0.0, 1.0] Residual algorithm parameter that
interpolates between Q-learning and residual gradient algorithm.
http://www.leemon.com/papers/1995b.pdf
debug_summaries: If True, add summaries to help debug behavior.
Raises:
ValueError: If 'dqda_clipping' is < 0.
"""
self._observation_spec = observation_spec[0]
self._action_spec = action_spec[0]
self._state_shape = tf.TensorShape([None]).concatenate(
self._observation_spec.shape)
self._action_shape = tf.TensorShape([None]).concatenate(
self._action_spec.shape)
self._num_action_dims = self._action_spec.shape.num_elements()
self._scope = tf.get_variable_scope().name
self._actor_net = tf.make_template(
self.ACTOR_NET_SCOPE, actor_net, create_scope_now_=True)
self._critic_net = tf.make_template(
self.CRITIC_NET_SCOPE, critic_net, create_scope_now_=True)
self._target_actor_net = tf.make_template(
self.TARGET_ACTOR_NET_SCOPE, actor_net, create_scope_now_=True)
self._target_critic_net = tf.make_template(
self.TARGET_CRITIC_NET_SCOPE, critic_net, create_scope_now_=True)
self._td_errors_loss = td_errors_loss
if dqda_clipping < 0:
raise ValueError('dqda_clipping must be >= 0.')
self._dqda_clipping = dqda_clipping
self._actions_regularizer = actions_regularizer
self._target_q_clipping = target_q_clipping
self._residual_phi = residual_phi
self._debug_summaries = debug_summaries
