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