def client_eval(incoming_model_weights, dataset): """Returns local outputs after evaluting `model_weights` on `dataset`.""" with tf.init_scope(): model = model_fn() model_weights = model_utils.ModelWeights.from_model(model) tf.nest.map_structure(lambda v, t: v.assign(t), model_weights, incoming_model_weights) def reduce_fn(num_examples, batch): model_output = model.forward_pass(batch, training=False) if model_output.num_examples is None: # Compute shape from the size of the predictions if model didn't use the # batch size. return num_examples + tf.shape(model_output.predictions, out_type=tf.int64)[0] else: return num_examples + tf.cast(model_output.num_examples, tf.int64) dataset_reduce_fn = dataset_reduce.build_dataset_reduce_fn( use_experimental_simulation_loop) num_examples = dataset_reduce_fn( reduce_fn=reduce_fn, dataset=dataset, initial_state_fn=lambda: tf.zeros([], dtype=tf.int64)) return collections.OrderedDict( local_outputs=model.report_local_outputs(), num_examples=num_examples)
def build(self, input_shape): if tf.__version__.startswith('1.5'): self.mean = self.add_weight(name='mean', shape=(self.b, input_shape[1]), initializer=tf.constant_initializer( np.tile(self.m, (self.b, 1)), verify_shape=True), trainable=False) self.std = self.add_weight(name='std', shape=(self.b, input_shape[1]), initializer=tf.constant_initializer( np.tile(self.s, (self.b, 1))), trainable=False) elif tf.__version__.startswith('2.'): with tf.init_scope(): self.mean = self.add_weight( name='mean', shape=(self.b, input_shape[1]), #initializer=tf.constant_initializer(np.tile(self.m,(self.b,1)),verify_shape=True), initializer=tf.constant_initializer( np.tile(self.m, (self.b, 1))), trainable=False) self.std = self.add_weight( name='std', shape=(self.b, input_shape[1]), #initializer=tf.constant_initializer(np.tile(self.s,(self.b,1)),verify_shape=True), initializer=tf.constant_initializer( np.tile(self.s, (self.b, 1))), trainable=False) else: sys.exit("Tensforflow version " + tf.__version__ + " unknown for preprocessing layer") super(PreprocessLayer, self).build(input_shape) # Be sure to call this at the end
def __call__(self, model, data): x, y = data batch_size = tf.shape(x)[0] with tf.GradientTape() as tape: y_pred = model(x, training=True) loss = model.compiled_loss(y, y_pred, regularization_losses=model.losses) gradients = tape.gradient(loss, model.trainable_weights) gradients = [g * tf.cast(batch_size, tf.float32) for g in gradients] sample = self.normal_initializer(tf.shape(y_pred))# * tf.sqrt(tf.cast(y_pred, tf.float32)) sample = model.sample_observations(y_pred, sample) sample = tf.reduce_sum(sample, 0, keepdims=True) single_x = tf.expand_dims(x[0,], axis=0) with tf.GradientTape() as tape: observations = model(single_x, training=True) fim_samples = tape.gradient(observations, model.trainable_weights, output_gradients=sample) fim_samples = [tf.expand_dims(samples, 0) for samples in fim_samples] fim_diagonal = model.fim_diagonal(x) with tf.init_scope(): model.optimizer._create_slots(model.trainable_weight_initializers()) model.optimizer.apply_information(zip(fim_samples, fim_diagonal, model.trainable_weights)) model.optimizer.apply_gradients(zip(gradients, model.trainable_weights)) model.compiled_metrics.update_state(y, y_pred) # Return a dict mapping metric names to current value return {m.name: m.result() for m in model.metrics}
def server_update(global_model, mean_model_delta, optimizer_state): """Updates the global model with the mean model update from clients.""" with tf.init_scope(): # Create a structure of variables that the server optimizer can update. model_variables = tf.nest.map_structure( lambda t: tf.Variable(initial_value=tf.zeros(t.shape, t.dtype) ), global_model) optimizer = keras_optimizer.build_or_verify_tff_optimizer( server_optimizer_fn, model_variables.trainable, disjoint_init_and_next=True) # Set the variables to the current global model, the optimizer will # update these variables. tf.nest.map_structure(lambda a, b: a.assign(b), model_variables, global_model) # We might have a NaN value e.g. if all of the clients processed had no # data, so the denominator in the federated_mean is zero. If we see any # NaNs, zero out the whole update. # TODO(b/124538167): We should increment a server counter to # track the fact a non-finite weights_delta was encountered. finite_weights_delta, _ = tensor_utils.zero_all_if_any_non_finite( mean_model_delta) # Update the global model variables with the delta as a pseudo-gradient. negative_weights_delta = tf.nest.map_structure(lambda w: -1.0 * w, finite_weights_delta) optimizer_state, updated_weights = optimizer.next( optimizer_state, model_variables.trainable, negative_weights_delta) # Keras optimizers mutate model variables in with the `next` step above, so # we skip calling the assignment for those optimizers. if not isinstance(optimizer, keras_optimizer.KerasOptimizer): tf.nest.map_structure(lambda a, b: a.assign(b), model_variables.trainable, updated_weights) return model_variables, optimizer_state
def call(self, inputs, training=None, **kwargs): with tf.init_scope(): if not self._is_set: self.layer(inputs) self._is_set = True if training: w = self.layer.kernel w_shape = w.shape.as_list() w = tf.reshape(w, (-1, self.d)) u_hat = self.u v_hat = None for _ in range(self.power_iteration): v_ = tf.matmul(u_hat, w, transpose_b=True) v_hat = self.l2_normalize(v_) u_ = tf.matmul(v_hat, w) u_hat = self.l2_normalize(u_) # u_hat = tf.stop_gradient(u_hat) # v_hat = tf.stop_gradient(v_hat) sigma = tf.matmul(tf.matmul(v_hat, w), u_hat, transpose_b=True) sigma = tf.reshape(sigma, ()) self.u.assign(u_hat) self.sigma.assign(sigma) return self.layer(inputs, **kwargs) / (self.sigma + K.epsilon())
def update(dataset: tf.data.Dataset, state: State, message: server.Message, model_fn: Callable, optimizer_fn: Callable) -> Output: with tf.init_scope(): model = model_fn(pos_weight=state.client_pos_weight) optimizer = optimizer_fn() message.model.assign_weights_to(model) def training_fn(num_examples, batch): with tf.GradientTape() as tape: outputs = model.forward_pass(batch, training=True) optimizer.apply_gradients( zip(tape.gradient(outputs.loss, model.trainable_variables), model.trainable_variables)) return num_examples + outputs.num_examples client_weight = dataset.reduce(tf.constant(0, dtype=tf.int32), training_fn) weights_delta = tf.nest.map_structure(lambda a, b: a - b, model.trainable_variables, message.model.trainable) return Output(weights_delta=weights_delta, metrics=model.report_local_outputs(), client_weight=tf.cast(client_weight, dtype=tf.float32), client_state=State( client_index=state.client_index, client_pos_weight=state.client_pos_weight))
def set_optimizer_params(optimizer, param_list): # Set optimizer tracking parameters if len(param_list) > 0: with tf.name_scope(optimizer._name): with tf.init_scope(): optimizer._create_all_weights(param_list) return optimizer
def __call__(self, model, data): x, y = data batch_size = tf.shape(x)[0] with tf.GradientTape() as tape: y_pred = model(x, training=True) loss = model.compiled_loss(y, y_pred, regularization_losses=model.losses) gradients = tape.gradient(loss, model.trainable_weights) gradients = [g * tf.cast(batch_size, tf.float32) for g in gradients] sample = model.sample_observations(y_pred, self.normal_initializer(tf.shape(y_pred))) def sample_fim(x_and_sample): with tf.GradientTape() as tape: y_pred = model(tf.expand_dims(x_and_sample[0], 0), training=True) return tape.gradient(y_pred, model.trainable_weights, output_gradients=tf.expand_dims(x_and_sample[1], 0)) fim_samples = tf.vectorized_map(sample_fim, (x, sample)) fim_diagonal = model.fim_diagonal(x) with tf.init_scope(): model.optimizer._create_slots(model.trainable_weight_initializers()) model.optimizer.apply_information(zip(fim_samples, fim_diagonal, model.trainable_weights)) model.optimizer.apply_gradients(zip(gradients, model.trainable_weights)) model.compiled_metrics.update_state(y, y_pred) # Return a dict mapping metric names to current value return {m.name: m.result() for m in model.metrics}
def _create_special_tokens_dict(self, vocab_table, vocab_file): special_tokens = dict(start_of_sequence_id="[CLS]", end_of_segment_id="[SEP]", padding_id="[PAD]", mask_id="[MASK]") with tf.init_scope(): if tf.executing_eagerly(): special_token_ids = vocab_table.lookup( tf.constant(list(special_tokens.values()), tf.string)) else: # A blast from the past: non-eager init context while building Model. # This can happen with Estimator or tf.compat.v1.disable_v2_behavior(). logging.warning( "Non-eager init context; computing " "BertTokenizer's special_tokens_dict in tf.compat.v1.Session" ) with tf.Graph().as_default(): local_vocab_table, _ = self._create_vocab_table_and_initializer( vocab_file) special_token_ids_tensor = local_vocab_table.lookup( tf.constant(list(special_tokens.values()), tf.string)) init_ops = [tf.compat.v1.initialize_all_tables()] with tf.compat.v1.Session() as sess: sess.run(init_ops) special_token_ids = sess.run(special_token_ids_tensor) result = dict() for k, v in zip(special_tokens, special_token_ids): v = int(v) # Numpy to Python. if v >= 0: result[k] = v else: logging.warning( "Could not find %s as token \"%s\" in vocab file %s", k, special_tokens[k], vocab_file) return result
def __init__( self, values_shape=(), values_dtype="int8", name="flip_ratio", dtype=None ): super().__init__(name=name, dtype=dtype) self.values_dtype = tf.as_dtype(values_dtype) self.values_shape = tf.TensorShape(values_shape).as_list() self.is_weight_metric = True with tf.init_scope(): self._previous_values = self.add_weight( "previous_values", shape=values_shape, dtype=self.values_dtype, initializer=tf.keras.initializers.zeros, aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA, ) self.total = self.add_weight( "total", initializer=tf.keras.initializers.zeros, aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA, ) self.count = self.add_weight( "count", initializer=tf.keras.initializers.zeros, aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA, ) self._size = np.prod(self.values_shape)
def get_summary_writer(save_dir, subdir='', comm=MPI.COMM_WORLD): if comm.Get_rank() != 0: return None if save_dir is None: return None with tf.init_scope(): return summary.create_file_writer(os.path.join(save_dir, 'tb', subdir))
def local_update(initial_weights, data): # TODO(b/190334722): Restructure so that model_fn only needs to be invoked # once. with tf.init_scope(): model = model_fn() model_weights = model_utils.ModelWeights.from_model(model) tf.nest.map_structure(lambda weight, value: weight.assign(value), model_weights, initial_weights) num_examples = tf.constant(0, tf.int32) optimizer_state = optimizer.initialize(optimizer_tensor_specs) # TODO(b/161529310): Different from creating an iterator using iter(data). for batch in data: with tf.GradientTape() as tape: outputs = model.forward_pass(batch) gradients = tape.gradient(outputs.loss, model_weights.trainable) num_examples += tf.shape(outputs.predictions)[0] optimizer_state, updated_weights = optimizer.next( optimizer_state, _flat_tuple(model_weights.trainable), _flat_tuple(gradients)) updated_weights = tf.nest.pack_sequence_as(model_weights.trainable, updated_weights) tf.nest.map_structure(lambda weight, value: weight.assign(value), model_weights.trainable, updated_weights) model_delta = tf.nest.map_structure(lambda x, y: x - y, initial_weights.trainable, model_weights.trainable) return ClientResult( update=model_delta, update_weight=tf.cast(num_examples, tf.float32))
def _build_from_signature(self, query, value, key=None): super(MultiHeadRelativeAttention, self)._build_from_signature(query=query, value=value, key=key) if hasattr(value, "shape"): value_shape = tf.TensorShape(value.shape) else: value_shape = value if key is None: key_shape = value_shape elif hasattr(key, "shape"): key_shape = tf.TensorShape(key.shape) else: key_shape = key common_kwargs = dict(kernel_initializer=self._kernel_initializer, bias_initializer=self._bias_initializer, kernel_regularizer=self._kernel_regularizer, bias_regularizer=self._bias_regularizer, activity_regularizer=self._activity_regularizer, kernel_constraint=self._kernel_constraint, bias_constraint=self._bias_constraint) with tf.init_scope(): einsum_equation, _, output_rank = _build_proj_equation( key_shape.rank - 1, bound_dims=1, output_dims=2) self._encoding_dense = tf.keras.layers.EinsumDense( einsum_equation, output_shape=_get_output_shape( output_rank - 1, [self._num_heads, self._key_dim]), bias_axes=None, name="encoding", **common_kwargs)
def _set_initializers(self): """Change initializers to load a model from a tensorflow checkpoint.""" if self.comm.Get_rank() > 0 or self.train_dir == 'test': return assert self.model.built checkpoint_scope = 'reward_model' with tf.init_scope(): # Initialize! params = {v.op.name: v for v in self.get_params()} checkpoint = tf.train.latest_checkpoint( os.path.join(self.train_dir, 'checkpoints/')) available = tf.train.list_variables(checkpoint) unchanged = {} for name, shape in available: if not name.startswith(checkpoint_scope + '/'): # print('skipping', name) continue if name.endswith('adam') or name.endswith('adam_1'): # print('skipping', name) continue print('setting', name) var = params[self.scope + name[len(checkpoint_scope):]] assert var.shape == shape, 'Shape mismatch: %s.shape = %s != %s' % ( var.op.name, var.shape, shape) unchanged[name] = var tf.train.init_from_checkpoint(checkpoint, unchanged)
def server_update(global_model, mean_model_delta, optimizer_state): """Updates the global model with the mean model update from clients.""" with tf.init_scope(): model = model_fn() optimizer = server_optimizer_fn() # We must force variable creation for momentum and adaptive optimizers. _eagerly_create_optimizer_variables(model=model, optimizer=optimizer) model_variables = model_utils.ModelWeights.from_model(model) optimizer_variables = optimizer.variables() # Set the variables to the current global model, the optimizer will # update these variables. tf.nest.map_structure(lambda a, b: a.assign(b), (model_variables, optimizer_variables), (global_model, optimizer_state)) # We might have a NaN value e.g. if all of the clients processed had no # data, so the denominator in the federated_mean is zero. If we see any # NaNs, zero out the whole update. # TODO(b/124538167): We should increment a server counter to # track the fact a non-finite weights_delta was encountered. finite_weights_delta, _ = tensor_utils.zero_all_if_any_non_finite( mean_model_delta) # Update the global model variables with the delta as a pseudo-gradient. _apply_delta(optimizer=optimizer, model=model, delta=finite_weights_delta) return model_variables, optimizer_variables
def __init__(self, spec, meta_graph, trainable, checkpoint_path, name): """Private constructor. Args: spec: _ModuleSpec instance. meta_graph: MetaGraphDef to use trainable: whether module is trainable. checkpoint_path: None or a string to the variables checkpoints. name: variable and scope name where to instantiate the Module. Must be an unused name scope. """ self._spec = spec self._meta_graph = meta_graph self._trainable = trainable self._checkpoint_path = checkpoint_path register_ops_if_needed({ op.name for op in self._meta_graph.meta_info_def.stripped_op_list.op }) # Use an init scope to clear dependencies and lift state ops outside of # function-building graphs. Modules can be constructed from deep inside # functions that have dependencies active. Note that the dependencies # would be active when applying the Module signature, just not active # when creating the Module state. This use case has showed up in some # TPU training code. with tf.init_scope(): self._init_state(name)
def _compute_rolled_dates_table(self, roll_convention): """Computes and caches rolled dates table.""" already_computed = self._table_cache.rolled_dates.get( roll_convention, None) if already_computed is not None: return already_computed roll_convention_np = _to_np_roll_convention(roll_convention) holidays_arg = self._holidays_np if holidays_arg is None: holidays_arg = [] # np.busday_offset doesn't accept None adjusted_np = np.busday_offset(dates=self._dates_np, offsets=0, roll=roll_convention_np, weekmask=1 - self._weekend_mask, holidays=holidays_arg) rolled_date_table = adjusted_np.astype(np.int32) + _ORDINAL_OF_1_1_1970 # To make tensor caching safe, lift the ops out of the current scope using # tf.init_scope(). This allows e.g. to cache these tensors in one # tf.function and reuse them in another tf.function. with tf.init_scope(): rolled_date_table = tf.convert_to_tensor(rolled_date_table, name="rolled_date_table") self._table_cache.rolled_dates[roll_convention] = rolled_date_table return rolled_date_table
def _get_temporary_graph_tensor(tensor_info, name): """Get a temporary graph tensor using attributes in `tensor_info`.""" is_asset_filepath = tensor_info.temporary_asset_value is not None if is_asset_filepath: # Placeholders cannot be used for assets as they will be initialized as part # of the init op. Hence, a temporary file is written out during tracing. # TODO(b/164921571) Support temporary files in tfrecord format. # TODO(b/149997088): Reduce number of temporary files written out. with tf.init_scope(): temporary_asset_filepath = os.path.join( TFGraphContext.get_or_create_temp_dir(), uuid.uuid4().hex) with tf.io.gfile.GFile(temporary_asset_filepath, 'w') as f: f.write(tensor_info.temporary_asset_value) result = tf.constant( temporary_asset_filepath, dtype=tensor_info.dtype, shape=tensor_info.shape, name=name) else: # Using a placeholder with no default value causes tracing to fail if there # is any control flow dependent on a child tensor of this placeholder. # Hence, provide a temporary default value for it. # If dtype is string, we want a tensor that contains '0's instead of b'[] to # allow string to numeric conversion ops to trace successfully. temporary_dtype = ( tf.int64 if tensor_info.dtype == tf.string else tensor_info.dtype) temporary_tensor = tf2_utils.supply_missing_tensor( 1, tf.TensorShape(tensor_info.shape), temporary_dtype) if tensor_info.dtype == tf.string: temporary_tensor = tf.strings.as_string(temporary_tensor) result = tf.raw_ops.PlaceholderWithDefault( input=temporary_tensor, shape=tensor_info.shape) return result
def is_xla_available(): """Is XLA compilation available for the current device context?""" global _IS_XLA_AVAILABLE # There's unfortunately no cleaner way to get the device other than creating a # new op and querying it. with tf.name_scope("is_xla_available"): device = tf.constant(0.0).device if device not in _IS_XLA_AVAILABLE: try: # Take ourselves outside of any tf.function calls. with tf.init_scope(): # Create temporary xla subgraph with tf.compat.v1.Graph().as_default(): # We'll use a session so we can be compatible with both TF1 and TF2 with tf.compat.v1.Session() as sess: # Check for XLA on the given device. with tf.device(device): sess.run( tf.xla.experimental.compile( lambda: tf.constant(0.0))) except (ValueError, tf.errors.InvalidArgumentError): _IS_XLA_AVAILABLE[device] = False else: _IS_XLA_AVAILABLE[device] = True return _IS_XLA_AVAILABLE[device]
def _get_temporary_analyzer_output( temp_dir: str, tensor_info: TensorInfo, name: Optional[str] = None) -> TemporaryAnalyzerOutputWrapper: """Create a temporary graph tensor using attributes in `tensor_info`. Args: temp_dir: Path to a directory to write out any temporary asset files to. tensor_info: A `TensorInfo` object containing attributes to create the graph tensor. name: A string (or None). The created graph tensor uses this name. Returns: A named tuple `TemporaryAnalyzerOutputWrapper` with: asset: If the graph tensor represents a path to an asset file, a `tf.saved_model.Asset` object for tracking. Else, None. graph_tensor: The graph tensor """ asset = None with tf.name_scope('temporary_analyzer_output'): is_asset_filepath = tensor_info.temporary_asset_value is not None if is_asset_filepath: # Placeholders cannot be used for assets, if this graph will be serialized # to a SavedModel, as they will be initialized with the init op. If a # `temp_dir` is provided, it is assumed that this graph will be # serialized and a temporary asset file is written out . Else, a # placeholder is returned. # TODO(b/164921571) Support temporary files in tfrecord format. # TODO(b/149997088): Reduce number of temporary files written out. if temp_dir: with tf.init_scope(): temporary_asset_filepath = os.path.join(temp_dir, uuid.uuid4().hex) with tf.io.gfile.GFile(temporary_asset_filepath, 'w') as f: f.write(tensor_info.temporary_asset_value) # Wrap asset files using `tf.saved_model.Asset` to ensure that # `SavedModel`s exported are hermetic. asset = common_types.Asset(temporary_asset_filepath) graph_tensor = tf.constant( temporary_asset_filepath, dtype=tensor_info.dtype, shape=tensor_info.shape, name=name) else: graph_tensor = tf.raw_ops.Placeholder( dtype=tensor_info.dtype, shape=tensor_info.shape, name=name) else: # Using a placeholder with no default value causes tracing to fail if # there is any control flow dependent on a child tensor of this # placeholder. Hence, provide a temporary default value for it. # If dtype is string, we want a tensor that contains '0's instead of b'[] # to allow string to numeric conversion ops to trace successfully. temporary_dtype = ( tf.int64 if tensor_info.dtype == tf.string else tensor_info.dtype) temporary_tensor = tf2_utils.supply_missing_tensor( 1, tf.TensorShape(tensor_info.shape), temporary_dtype) if tensor_info.dtype == tf.string: temporary_tensor = tf.strings.as_string(temporary_tensor) graph_tensor = tf.raw_ops.PlaceholderWithDefault( input=temporary_tensor, shape=tensor_info.shape, name=name) return TemporaryAnalyzerOutputWrapper(asset, graph_tensor)
def _check_same_graph(self): """Checks that the module is not being connect to multiple Graphs. An instance of a Sonnet module 'owns' the variables it contains, and permits seamless variable sharing. As such, connecting a single module instance to multiple Graphs is not possible - this function will raise an error should that occur. Raises: DifferentGraphError: if the module is connected to a different Graph than it was previously used in. """ with tf.init_scope(): # We need `init_scope` incase we're running inside a defun. In that case # what we want is information about where the function will be called not # where the function is being built. current_graph = tf.get_default_graph() will_call_in_eager_context = tf.executing_eagerly() if self._graph is None: self._graph = current_graph self._set_module_info() if not will_call_in_eager_context: # Same graph checks only make sense when calling from graph mode (in eager # mode there is a single process level context where all modules are # created). if self._graph != current_graph: raise DifferentGraphError( "Cannot connect module to multiple Graphs.")
def apply_gradients(self, grads_and_vars, name: Optional[str] = None, **kwargs): """Apply gradients to variables for each optimizer. On the first call to `apply_gradients()`, compute the mapping from variables to optimizers and cache it in the `self.var_opt_mapping` dict for serialization and faster access. """ if self.var_opt_mapping is None: # Convert `grads_and_vars` to list so we can iterate multiple times over it grads_and_vars = list(grads_and_vars) self._compute_var_opt_mapping(grads_and_vars) # Split gradients and variables into a separate list for each optimizer grad_var_lists = [[] for _ in range(len(self.pred_opt_pairs) + 1)] for grad, var in grads_and_vars: if var.name in self.var_opt_mapping: grad_var_lists[self.var_opt_mapping[var.name]].append( (grad, var)) with tf.init_scope(): for optimizer, opt_grads_and_vars in zip(self.optimizers, grad_var_lists): optimizer._create_slots([v for (_, v) in grads_and_vars]) return tf.distribute.get_replica_context().merge_call( self._apply_gradients, args=(grad_var_lists, name), kwargs=kwargs)
def _get_step(self): with tf.init_scope(): if context.executing_eagerly(): graph = None else: graph = tf.get_default_graph() return self._get_non_slot_variable('step', graph=graph)
def finalizer(unfinalized_metric_values: List[tf.Tensor]): # Construct a new keras metirc here, which is necessary because this # `tf.function` may be invoked in a different context as the `model_fn`, and # we need the `tf.Variable`s to be created in the current scope in order to # use `keras_metric.result()`. with tf.init_scope(): keras_metric = create_keras_metric(metric) py_typecheck.check_type(unfinalized_metric_values, list) if len(keras_metric.variables) != len(unfinalized_metric_values): raise ValueError( 'The input to the finalizer should be a list of `tf.Tensor`s matching' f' the variables of the Keras metric {keras_metric.name}. Expected ' f'a list of `tf.Tensor`s of length {len(keras_metric.variables)}, ' f'found a list of length {len(unfinalized_metric_values)}.') for v, a in zip(keras_metric.variables, unfinalized_metric_values): py_typecheck.check_type(a, tf.Tensor) if v.shape != a.shape or v.dtype != a.dtype: raise ValueError( 'The input to the finalizer should be a list of `tf.Tensor`s ' f'matching the variables of the Keras metric {keras_metric.name}. ' f'Expected a `tf.Tensor` of shape {v.shape} and dtype {v.dtype!r}, ' f'found a `tf.Tensor` of shape {a.shape} and dtype {a.dtype!r}.') v.assign(a) return keras_metric.result()
def __init__( self, active=True, scale=1., print_loss=False, print_batch_time=False, return_lossval=False, print_time=False, #compat, has no effect **kwargs): super(LossLayerBase, self).__init__(**kwargs) if print_time: print( "print_time has no effect and is only for compatibility purposes" ) self.active = active self.scale = scale self.print_loss = print_loss self.print_batch_time = print_batch_time self.return_lossval = return_lossval with tf.init_scope(): now = tf.timestamp() self.time = tf.Variable(-now, name=self.name + '_time', trainable=False)
def evaluate(dataset: tf.data.Dataset, state: State, weights: tff.learning.ModelWeights, coefficient_fn: Callable, model_fn: Callable) -> Evaluation: with tf.init_scope(): mixing_coefficients = coefficient_fn() model = model_fn(pos_weight=state.client_pos_weight) tf.nest.map_structure(lambda v, t: v.assign(t), mixing_coefficients, state.mixing_coefficients) __mix_weights(mixing_coefficients, state.model, weights).assign_weights_to(model) def evaluation_fn(state, batch): outputs = model.forward_pass(batch, training=False) y_true = tf.reshape(batch[1], (-1, )) y_pred = tf.round( tf.nn.sigmoid(tf.reshape(outputs.predictions, (-1, )))) return state + tf.math.confusion_matrix(y_true, y_pred, num_classes=2) confusion_matrix = dataset.reduce(tf.zeros((2, 2), dtype=tf.int32), evaluation_fn) return Evaluation(confusion_matrix=confusion_matrix, metrics=model.report_local_outputs())
def initialize(self, input_shape): self.input_shape = input_shape self.Dx_loss, self.Gx_loss = self.adv_loss_ctor(self.D_x) self.Dy_loss, self.Gy_loss = self.adv_loss_ctor(self.D_y) with tf.init_scope(): self.G_zx.initialize(input_shape) self.G_zy.initialize(input_shape) self._init_checkpoint()
def __init__(self, initializer): super(Index, self).__init__() if isinstance(initializer, trackable_base.Trackable): self._initializer = self._track_trackable(initializer, "_initializer") with tf.init_scope(): self._resource_handle = self._create_resource() self._init_op = self._initialize()
def _generate(self, feature_map_shape_list): """Generates a collection of bounding boxes to be used as anchors. Args: feature_map_shape_list: list of pairs of convnet layer resolutions in the format [(height_0, width_0)]. For example, setting feature_map_shape_list=[(8, 8)] asks for anchors that correspond to an 8x8 layer. For this anchor generator, only lists of length 1 are allowed. Returns: boxes_list: a list of BoxLists each holding anchor boxes corresponding to the input feature map shapes. Raises: ValueError: if feature_map_shape_list, box_specs_list do not have the same length. ValueError: if feature_map_shape_list does not consist of pairs of integers """ if not (isinstance(feature_map_shape_list, list) and len(feature_map_shape_list) == 1): raise ValueError( 'feature_map_shape_list must be a list of length 1.') if not all([ isinstance(list_item, tuple) and len(list_item) == 2 for list_item in feature_map_shape_list ]): raise ValueError('feature_map_shape_list must be a list of pairs.') # Create constants in init_scope so they can be created in tf.functions # and accessed from outside of the function. with tf.init_scope(): self._base_anchor_size = tf.cast(tf.convert_to_tensor( self._base_anchor_size), dtype=tf.float32) self._anchor_stride = tf.cast(tf.convert_to_tensor( self._anchor_stride), dtype=tf.float32) self._anchor_offset = tf.cast(tf.convert_to_tensor( self._anchor_offset), dtype=tf.float32) grid_height, grid_width = feature_map_shape_list[0] scales_grid, aspect_ratios_grid = ops.meshgrid(self._scales, self._aspect_ratios) scales_grid = tf.reshape(scales_grid, [-1]) aspect_ratios_grid = tf.reshape(aspect_ratios_grid, [-1]) anchors = tile_anchors(grid_height, grid_width, scales_grid, aspect_ratios_grid, self._base_anchor_size, self._anchor_stride, self._anchor_offset) num_anchors = anchors.num_boxes_static() if num_anchors is None: num_anchors = anchors.num_boxes() anchor_indices = tf.zeros([num_anchors]) anchors.add_field('feature_map_index', anchor_indices) return [anchors]
def _create_optimizer(self): """Initializes the hyperparameters and sets the self._optimizer property.""" if self._optimizer: return if not self._layer_collection: self.register_layers(self._model, self._loss) if self._config['adapt_damping']: if 'train_batch' not in self._kfac_kwargs: raise ValueError( 'Must provide a train_batch tuple to use adaptive ' 'damping. Use register_train_batch or pass it in ' 'during optimizer construction.') if 'loss_fn' not in self._kfac_kwargs: self._kfac_kwargs['loss_fn'] = utils.get_loss_fn( self._model, self._loss, loss_weights=self._config['loss_weights']) with tf.name_scope(self._name): with tf.init_scope(): # "iterations" property will create iterations if necessary. _ = self.iterations self._create_hypers() self._kfac_kwargs.update(self._hyper) try: # We use the TF 1 variable_scope instead of the TF 2 recommended # name_scope because we need to recover the variables created in this # scope, which is not possible with name_scope. with tf.variable_scope(self._tf_var_scope): self._optimizer = _KFAC_OPT_CLASS( layer_collection=self._layer_collection, **self._kfac_kwargs) except ValueError as e: msg = str(e) if re.search('Variable .* already exists', msg): raise ValueError( 'You may have instantiated a KFAC Optimizer with the same name as ' 'an existing one. Try resetting the default graph, instantiating ' 'the optimizer with a different name, or changing the optimizer\'s ' 'name.\nHere is the original ValueError:\n ' + msg) elif re.search( 'Found the following errors with variable registration' '.*gamma.*registered with wrong number of uses.*', msg): # We don't regex the name batch_normalization because the user could # have renamed the layer. We don't regex beta because they could have # used BatchNorm without the shift. raise ValueError( 'There may have been an issue registering BatchNormalization. Try ' 'using tf.keras.backend.set_learning_phase before model ' 'construction. An alternative solution is to use the unfused ' 'batchnorm implementation (pass the argument fused=False to ' 'BatchNormalization).\nHere is the original ValueError:\n ' + msg) else: raise e