Python tensorflow.TensorSpec方法代码示例

# 需要导入模块: import tensorflow [as 别名]
# 或者: from tensorflow import TensorSpec [as 别名]
def as_simple_encoder(encoder, tensorspec):
  """Wraps an `Encoder` object as a `SimpleEncoder`.

  Args:
    encoder: An `Encoder` object to be used to encoding.
    tensorspec: A `TensorSpec`. The created `SimpleEncoder` will be constrained
      to only encode input values compatible with `tensorspec`.

  Returns:
    A `SimpleEncoder`.

  Raises:
    TypeError:
      If `encoder` is not an `Encoder` or `tensorspec` is not a `TensorSpec`.
  """
  if not isinstance(encoder, core_encoder.Encoder):
    raise TypeError('The encoder must be an instance of `Encoder`.')
  if not isinstance(tensorspec, tf.TensorSpec):
    raise TypeError('The tensorspec must be a tf.TensorSpec.')
  return simple_encoder.SimpleEncoder(encoder, tensorspec) 

# 需要导入模块: import tensorflow [as 别名]
# 或者: from tensorflow import TensorSpec [as 别名]
def assert_compatible(spec, value):
  """Asserts that values are compatible with given specs.

  Args:
    spec: A structure compatible with `tf.nest`, with `tf.TensorSpec` values.
    value: A collection of values that should be compatible with `spec`. Must be
      the same structure as `spec`.

  Raises:
    TypeError: If `spec` does not contain only `tf.TensorSpec` objects.
    ValueError: If the provided `value` is not compatible with `spec`.
  """

  def validate_spec(s, v):
    if not isinstance(s, tf.TensorSpec):
      raise TypeError('Each value in `spec` must be a tf.TensorSpec.')
    return s.is_compatible_with(v)

  compatible = tf.nest.map_structure(validate_spec, spec, value)
  if not all(tf.nest.flatten(compatible)):
    raise ValueError('The provided value is not compatible with spec.') 

# 需要导入模块: import tensorflow [as 别名]
# 或者: from tensorflow import TensorSpec [as 别名]
def test_none_state_equal_to_initial_state(self):
    """Tests that not providing state is the same as initial_state."""
    x = tf.constant(1.0)
    encoder = simple_encoder.SimpleEncoder(
        core_encoder.EncoderComposer(
            test_utils.PlusOneOverNEncodingStage()).make(),
        tf.TensorSpec.from_tensor(x))

    state = encoder.initial_state()
    stateful_iteration = _make_iteration_function(encoder)

    @tf.function
    def stateless_iteration(x):
      encoded_x, _ = encoder.encode(x)
      decoded_x = encoder.decode(encoded_x)
      return encoded_x, decoded_x

    _, encoded_x_stateful, decoded_x_stateful, _ = self.evaluate(
        stateful_iteration(x, state))
    encoded_x_stateless, decoded_x_stateless = self.evaluate(
        stateless_iteration(x))

    self.assertAllClose(encoded_x_stateful, encoded_x_stateless)
    self.assertAllClose(decoded_x_stateful, decoded_x_stateless) 

# 需要导入模块: import tensorflow [as 别名]
# 或者: from tensorflow import TensorSpec [as 别名]
def test_input_signature_enforced(self):
    """Tests that encode/decode input signature is enforced."""
    x = tf.constant(1.0)
    encoder = simple_encoder.SimpleEncoder(
        core_encoder.EncoderComposer(
            test_utils.PlusOneOverNEncodingStage()).make(),
        tf.TensorSpec.from_tensor(x))

    state = encoder.initial_state()
    with self.assertRaises(ValueError):
      bad_x = tf.stack([x, x])
      encoder.encode(bad_x, state)
    with self.assertRaises(ValueError):
      bad_state = state + (x,)
      encoder.encode(x, bad_state)
    encoded_x = encoder.encode(x, state)
    with self.assertRaises(ValueError):
      bad_encoded_x = dict(encoded_x)
      bad_encoded_x.update({'x': x})
      encoder.decode(bad_encoded_x) 

# 需要导入模块: import tensorflow [as 别名]
# 或者: from tensorflow import TensorSpec [as 别名]
def test_basic_encode_decode(self):
    """Tests basic encoding and decoding works as expected."""
    x_fn = lambda: tf.random.uniform((12,))
    encoder = gather_encoder.GatherEncoder.from_encoder(
        core_encoder.EncoderComposer(
            test_utils.PlusOneOverNEncodingStage()).make(),
        tf.TensorSpec.from_tensor(x_fn()))

    num_summands = 3
    iteration = _make_iteration_function(encoder, x_fn, num_summands)
    state = encoder.initial_state()

    for i in range(1, 5):
      data = self.evaluate(iteration(state))
      for j in range(num_summands):
        self.assertAllClose(
            data.x[j] + 1 / i,
            _encoded_x_field(data.encoded_x[j], [TENSORS, PN_VALS]))
      self.assertEqual((i,), data.initial_state)
      self.assertEqual((i + 1,), data.updated_state)
      state = data.updated_state 

# 需要导入模块: import tensorflow [as 别名]
# 或者: from tensorflow import TensorSpec [as 别名]
def test_none_state_equal_to_initial_state(self):
    """Tests that not providing state is the same as initial_state."""
    x_fn = lambda: tf.constant(1.0)
    encoder = gather_encoder.GatherEncoder.from_encoder(
        core_encoder.EncoderComposer(
            test_utils.PlusOneOverNEncodingStage()).make(),
        tf.TensorSpec.from_tensor(x_fn()))

    num_summands = 3
    stateful_iteration = _make_iteration_function(encoder, x_fn, num_summands)
    state = encoder.initial_state()
    stateless_iteration = _make_stateless_iteration_function(
        encoder, x_fn, num_summands)

    stateful_data = self.evaluate(stateful_iteration(state))
    stateless_data = self.evaluate(stateless_iteration())

    self.assertAllClose(stateful_data.encoded_x, stateless_data.encoded_x)
    self.assertAllClose(stateful_data.decoded_x, stateless_data.decoded_x) 

# 需要导入模块: import tensorflow [as 别名]
# 或者: from tensorflow import TensorSpec [as 别名]
def test_commutativity_with_sum(self):
    """Tests that encoder that commutes with sum works."""
    x_fn = lambda: tf.constant([1.0, 3.0])
    encoder = gather_encoder.GatherEncoder.from_encoder(
        core_encoder.EncoderComposer(test_utils.TimesTwoEncodingStage()).make(),
        tf.TensorSpec.from_tensor(x_fn()))

    for num_summands in [1, 3, 7]:
      iteration = _make_iteration_function(encoder, x_fn, num_summands)
      data = self.evaluate(iteration(encoder.initial_state()))
      for i in range(num_summands):
        self.assertAllClose([1.0, 3.0], data.x[i])
        self.assertAllClose(
            [2.0, 6.0], _encoded_x_field(data.encoded_x[i], [TENSORS, T2_VALS]))
        self.assertAllClose(list(data.part_decoded_x[i].values())[0],
                            list(data.encoded_x[i].values())[0])
      self.assertAllClose(np.array([2.0, 6.0]) * num_summands,
                          list(data.summed_part_decoded_x.values())[0])
      self.assertAllClose(np.array([1.0, 3.0]) * num_summands, data.decoded_x) 

# 需要导入模块: import tensorflow [as 别名]
# 或者: from tensorflow import TensorSpec [as 别名]
def test_full_commutativity_with_sum(self):
    """Tests that fully commutes with sum property works."""
    spec = tf.TensorSpec((2,), tf.float32)

    encoder = gather_encoder.GatherEncoder.from_encoder(
        core_encoder.EncoderComposer(test_utils.TimesTwoEncodingStage()).make(),
        spec)
    self.assertTrue(encoder.fully_commutes_with_sum)

    encoder = gather_encoder.GatherEncoder.from_encoder(
        core_encoder.EncoderComposer(
            test_utils.TimesTwoEncodingStage()).add_parent(
                test_utils.TimesTwoEncodingStage(), T2_VALS).make(), spec)
    self.assertTrue(encoder.fully_commutes_with_sum)

    encoder = core_encoder.EncoderComposer(
        test_utils.SignIntFloatEncodingStage())
    encoder.add_child(test_utils.TimesTwoEncodingStage(), SIF_SIGNS)
    encoder.add_child(test_utils.PlusOneEncodingStage(), SIF_INTS)
    encoder.add_child(test_utils.TimesTwoEncodingStage(), SIF_FLOATS).add_child(
        test_utils.PlusOneOverNEncodingStage(), T2_VALS)
    encoder = gather_encoder.GatherEncoder.from_encoder(encoder.make(), spec)
    self.assertFalse(encoder.fully_commutes_with_sum) 

# 需要导入模块: import tensorflow [as 别名]
# 或者: from tensorflow import TensorSpec [as 别名]
def test_synthetic(self):
    client_data = emnist.get_synthetic(num_clients=4)
    self.assertLen(client_data.client_ids, 4)

    self.assertEqual(
        client_data.element_type_structure,
        collections.OrderedDict([
            ('pixels', tf.TensorSpec(shape=(28, 28), dtype=tf.float32)),
            ('label', tf.TensorSpec(shape=(), dtype=tf.int32)),
        ]))

    for client_id in client_data.client_ids:
      data = self.evaluate(
          list(client_data.create_tf_dataset_for_client(client_id)))
      images = [x['pixels'] for x in data]
      labels = [x['label'] for x in data]
      self.assertLen(labels, 10)
      self.assertCountEqual(labels, list(range(10)))
      self.assertLen(images, 10)
      self.assertEqual(images[0].shape, (28, 28))
      self.assertEqual(images[-1].shape, (28, 28)) 

# 需要导入模块: import tensorflow [as 别名]
# 或者: from tensorflow import TensorSpec [as 别名]
def tensor_spec_for_batch(dummy_batch):
  """Returns a TensorSpec for the given batch."""
  # TODO(b/131085687): Consider common util shared with model_utils.py.
  if hasattr(dummy_batch, '_asdict'):
    dummy_batch = dummy_batch._asdict()

  def _get_tensor_spec(tensor):
    # Convert input to tensors, possibly from nested lists that need to be
    # converted to a single top-level tensor.
    tensor = tf.convert_to_tensor(tensor)
    # Remove the batch dimension and leave it unspecified.
    spec = tf.TensorSpec(
        shape=[None] + tensor.shape.dims[1:], dtype=tensor.dtype)
    return spec

  return tf.nest.map_structure(_get_tensor_spec, dummy_batch)


# Set cmp=False to get a default hash function for tf.function. 

# 需要导入模块: import tensorflow [as 别名]
# 或者: from tensorflow import TensorSpec [as 别名]
def _broadcast_encoder_fn(value):
  """Function for building encoded broadcast.

  This method decides, based on the tensor size, whether to use lossy
  compression or keep it as is (use identity encoder). The motivation for this
  pattern is due to the fact that compression of small model weights can provide
  only negligible benefit, while at the same time, lossy compression of small
  weights usually results in larger impact on model's accuracy.

  Args:
    value: A tensor or variable to be encoded in server to client communication.

  Returns:
    A `te.core.SimpleEncoder`.
  """
  # TODO(b/131681951): We cannot use .from_tensor(...) because it does not
  # currently support Variables.
  spec = tf.TensorSpec(value.shape, value.dtype)
  if value.shape.num_elements() > 10000:
    return te.encoders.as_simple_encoder(
        te.encoders.uniform_quantization(FLAGS.broadcast_quantization_bits),
        spec)
  else:
    return te.encoders.as_simple_encoder(te.encoders.identity(), spec) 

# 需要导入模块: import tensorflow [as 别名]
# 或者: from tensorflow import TensorSpec [as 别名]
def inputs(self):
        return [tf.TensorSpec([None, self.image_shape, self.image_shape, 3], self.image_dtype, 'input'),
                tf.TensorSpec([None], tf.int32, 'label')] 

# 需要导入模块: import tensorflow [as 别名]
# 或者: from tensorflow import TensorSpec [as 别名]
def get_inputs(batch):
    return [tf.TensorSpec(
                (batch, 3, 32, 32) if DATA_FORMAT == "NCHW" else (batch, 32, 32, 3),
                tf.float32, 'input'),
            tf.TensorSpec((batch, 10), tf.float32, 'label')] 

# 需要导入模块: import tensorflow [as 别名]
# 或者: from tensorflow import TensorSpec [as 别名]
def inputs(self):
        return [tf.TensorSpec([args.batch, INPUT_SHAPE, INPUT_SHAPE, 3], IMAGE_DTYPE, 'input'),
                tf.TensorSpec([args.batch], tf.int32, 'label')] 

# 需要导入模块: import tensorflow [as 别名]
# 或者: from tensorflow import TensorSpec [as 别名]
def inputs(self):
        return [tf.TensorSpec([None, 3, 224, 224], tf.float32, 'input'),
                tf.TensorSpec([None], tf.int32, 'label')]