def export_saved_model(self, export_dir='.'): """export a saved model Args: export_dir: directory to save the saved model """ sess = tf.get_default_session() builder = tf.saved_model.builder.SavedModelBuilder(export_dir) tf.identity_n(self.outputs, name='output/hr') builder.add_meta_graph_and_variables(sess, tf.saved_model.tag_constants.SERVING) builder.save()
def decorated(*args, **kwargs): """Decorated function with custom gradient.""" if kwargs: raise ValueError( "The custom_gradient decorator currently suports keywords " "arguments only when eager execution is enabled.") name = "CustomGradient-%s" % ops.uid() args = [ops.convert_to_tensor(x) for x in args] result, grad_fn = f(*args) flat_result = nest.flatten(result) all_tensors = flat_result + args @ops.RegisterGradient(name) def internal_grad_fn(unused_op, *result_grads): gradients = nest.flatten(grad_fn(*result_grads[:len(flat_result)])) # Need to return one value per input to the IdentityN, so pad the # gradients of the inputs of the custom_gradient function with the # gradients of the outputs as well. return ([None] * len(flat_result)) + gradients with ops.get_default_graph().gradient_override_map({"IdentityN": name}): all_tensors = tf.identity_n(all_tensors) return nest.pack_sequence_as( structure=result, flat_sequence=all_tensors[:len(flat_result)])
def to_graph(self, graph): with graph.as_default(): placeholder = tf.keras.backend.placeholder( shape=(self.__num_inputs, ), name='MainInput', dtype=tf.float32) return tf.identity_n(tf.unstack(placeholder), name="Inp")
def build_graph(parameters): """Make a set of tests to do identity.""" input_tensors = [] input_count = (2 if parameters["op_to_use"] == "identity_n_with_2_inputs" else 1) input_tensors = [ tf.placeholder( dtype=tf.float32, name="input", shape=parameters["input_shape"]) for _ in range(input_count) ] # We add the Multiply before Identity just as a walk-around to make the test # pass when input_shape is scalar. # During graph transformation, TOCO will replace the Identity op with # Reshape when input has shape. However, currently TOCO can't distinguish # between missing shape and scalar shape. As a result, when input has scalar # shape, this conversion still fails. # TODO(b/129197312), remove the walk-around code once the bug is fixed. inputs_doubled = [input_tensor * 2.0 for input_tensor in input_tensors] if parameters["op_to_use"] == "identity": identity_outputs = [tf.identity(inputs_doubled[0])] elif parameters["op_to_use"] == "snapshot": identity_outputs = [array_ops.snapshot(inputs_doubled[0])] elif parameters["op_to_use"] in ("identity_n", "identity_n_with_2_inputs"): identity_outputs = tf.identity_n(inputs_doubled) return input_tensors, identity_outputs
def __init__(self, adj, x, labels, idx_train, idx_unlabeled=None, surrogate=None, surrogate_args={}, surrogate_lr=5e-3, seed=None, name=None, device='CPU:0', **kwargs): super().__init__(adj, x, labels, idx_train=idx_train, idx_unlabeled=idx_unlabeled, surrogate=surrogate, surrogate_args=surrogate_args, seed=seed, device=device, **kwargs) with tf.device(self.device): self.stored_weights = tf.identity_n(self.surrogate.weights) self.optimizer = Adam(surrogate_lr)
def _get_init_state_tuple(self, cell, name): _init_zero_states = cell.zero_state(self.hparams.batch_size, tf.float32) if self.stateful: _init_states = [tf.identity_n(_state, name) for _state in _init_zero_states] return tuple( [tf.nn.rnn_cell.LSTMStateTuple(*_state) for _state in _init_states] ) return _init_zero_states
def to_graph(self, graph, max_len): with graph.as_default(): placeholder = tf.keras.backend.placeholder(shape=(), dtype=tf.float32, name='sw') weights = [] for i in range(int(max_len) + 1): weights.append(tf.fill((i, ), placeholder)) return tf.identity_n(weights, "shared_weight")
def process(self, surrogate, train_nodes, unlabeled_nodes=None, lr=5e-3, reset=True): super().process(surrogate, train_nodes, unlabeled_nodes, reset=False) with tf.device(self.device): self.stored_weights = tf.identity_n(self.surrogate.weights) self.optimizer = Adam(lr) if reset: self.reset() return self
def export_model_pb(self, export_dir='.', export_name='model.pb', **kwargs): """ Export model as a constant protobuf. Unlike saved model, this one is not trainable Args: export_dir: directory to save the exported model export_name: model name """ self.output = tf.identity_n(self.output, name='output/hr') sess = tf.get_default_session() graph = sess.graph.as_graph_def() graph = tf.graph_util.remove_training_nodes(graph) graph = tf.graph_util.convert_variables_to_constants(sess, graph, [outp.name.split(':')[0] for outp in self.output]) tf.train.write_graph(graph, export_dir, export_name, as_text=False) tf.logging.info("Model exported to [%s/%s]." % (Path(export_dir).resolve(), export_name))
def init_normalization(x, *, name, init_scale=1., init, ema): with tf.variable_scope(name): g = get_var('g', shape=x.shape[1:], initializer=tf.constant_initializer(1.), ema=ema) b = get_var('b', shape=x.shape[1:], initializer=tf.constant_initializer(0.), ema=ema) if init: # data based normalization m_init, v_init = tf.nn.moments(x, [0]) scale_init = init_scale * tf.rsqrt(v_init + 1e-8) assert m_init.shape == v_init.shape == scale_init.shape == g.shape == b.shape with tf.control_dependencies([ g.assign(scale_init), b.assign(-m_init * scale_init) ]): g, b = tf.identity_n([g, b]) return g, b
def test_reversible_step(self, distribution): # Reversible layers satisfy: (a) strides = 1 (b) in_filter = out_filter bsz, h, w, c = 8, 32, 32, 32 filters = c strides = 1 input_tensor = tf.random.uniform(shape=[bsz, h, w, c]) with distribution.scope(): f = nn_blocks.ResidualInner( filters=filters // 2, strides=strides, batch_norm_first=False) g = nn_blocks.ResidualInner( filters=filters // 2, strides=1, batch_norm_first=False) test_layer = nn_blocks.ReversibleLayer(f, g) test_layer(input_tensor, training=False) # init weights optimizer = tf.keras.optimizers.SGD(learning_rate=0.01) @tf.function def step_fn(): with tf.GradientTape() as tape: output = test_layer(input_tensor, training=True) grads = tape.gradient(output, test_layer.trainable_variables) # Test applying gradients with optimizer works optimizer.apply_gradients(zip(grads, test_layer.trainable_variables)) return output @tf.function def fwd(): test_layer(input_tensor) distribution.run(fwd) # Initialize variables prev_variables = tf.identity_n(test_layer.trainable_variables) replica_output = distribution.run(step_fn) outputs = distribution.experimental_local_results(replica_output) # Assert variables values have changed values for v0, v1 in zip(prev_variables, test_layer.trainable_variables): self.assertNotAllEqual(v0, v1) # Assert forward pass shape expected_output_shape = [bsz, h // strides, w // strides, filters] for output in outputs: self.assertEqual(expected_output_shape, output.shape.as_list())
def _dky(x, y): with tf.GradientTape() as t: # Get the numbers of inputs. nx = B.shape(x)[0] ny = B.shape(y)[0] # Copy the input `nx` times to efficiently compute many derivatives. yis = tf.identity_n([y[:, i:i + 1]] * nx) t.watch(yis) # Tile inputs for batched computation. x = B.reshape(B.tile(x, 1, ny), nx * ny, -1) y = B.tile(y, nx, 1) # Insert tracked dimension, which is different for every tile. yi = B.concat(*yis, axis=0) y = B.concat(y[:, :i], yi, y[:, i + 1:], axis=1) # Perform the derivative computation. out = dense(k_elwise(x, y)) grads = t.gradient(out, yis, unconnected_gradients='zero') return B.transpose(B.concat(*grads, axis=1))
def export_freeze_model(self, export_dir='.', version=1): """export model as a constant protobuf. Unlike saved model, this one is not trainable Args: export_dir: directory to save the exported model version: version of the exported model """ self.outputs = tf.identity_n(self.outputs, name='output/hr') sess = tf.get_default_session() export_path = Path(export_dir) / str(version) while export_path.exists(): version += 1 # step ahead 1 version export_path = Path(export_dir) / str(version) export_path = str(export_path) graph = sess.graph.as_graph_def() graph = tf.graph_util.remove_training_nodes(graph) graph = tf.graph_util.convert_variables_to_constants( sess, graph, [outp.name.split(':')[0] for outp in self.outputs]) tf.train.write_graph(graph, export_path, self.name, as_text=False) tf.logging.info(f"Model exported to {export_path}/{self.name}.")
def _dkx(x, y): import tensorflow as tf with tf.GradientTape() as t: # Get the numbers of inputs. nx = B.shape(x)[0] ny = B.shape(y)[0] # Copy the input `ny` times to efficiently compute many derivatives. xis = tf.identity_n([x[:, i:i + 1]] * ny) t.watch(xis) # Tile inputs for batched computation. x = B.tile(x, ny, 1) y = B.reshape(B.tile(y, 1, nx), ny * nx, -1) # Insert tracked dimension, which is different for every tile. xi = B.concat(*xis, axis=0) x = B.concat(x[:, :i], xi, x[:, i + 1:], axis=1) # Perform the derivative computation. out = dense(k_elwise(x, y)) grads = t.gradient(out, xis, unconnected_gradients='zero') return B.concat(*grads, axis=1)
xentropy = tf.nn.sparse_softmax_cross_entropy_with_logits( labels=y, logits=logits) loss = tf.reduce_mean(xentropy, name="loss") learning_rate = tf.placeholder_with_default(0.01, shape=[], name="learning_rate") with tf.name_scope("train_simple"): optimizer = tf.train.GradientDescentOptimizer(learning_rate) training_op = optimizer.minimize(loss, name='optimize') with tf.name_scope("train"): optimizer = tf.train.GradientDescentOptimizer(learning_rate) gradients = optimizer.compute_gradients(loss) gradients_named = tf.identity_n(gradients, name='compute_gradients') gradients_named_length = tf.Variable( len(gradients_named), name='compute_gradients_output_length', trainable=False) training_op = optimizer.apply_gradients(gradients, name='apply_gradients') with tf.name_scope("eval"): correct = tf.nn.in_top_k(logits, y, 1) accuracy = tf.reduce_mean(tf.cast(correct, tf.float32), name="accuracy") init = tf.global_variables_initializer() with open('graph.pb', 'wb') as f:
def _get_identity_states(states, name): # use `tf.identity_n` function to access tensors by name _states = [tf.identity_n(_state, name) for _state in states] return tuple([tf.nn.rnn_cell.LSTMStateTuple(*_state) for _state in _states])
def sample_gamma_RSVI(z, alpha, fz): """Given a (frozen) Gamma sample z, and shape parameter alpha and log-prob fz, get RSVI gradient of sample. The key tensorflow dynamic is to overwrite the gradient of the identity_N op to work as the RSVI gradient.""" sample, _, _ = tf.identity_n([z, alpha, fz]) return sample
def model(self, m): # Back up if isinstance(m, tf.keras.Model) and m.weights: self.backup = tf.identity_n(m.weights) # TODO assert m is None or isinstance(m, tf.keras.Model) or torch.nn.Module self._model = m
def _common(cls, node, **kwargs): x = kwargs["tensor_dict"][node.inputs[0]] if isinstance(x, (list, tuple)): return [tf.identity_n(x)] else: return [tf.identity(x)]