Tipi di estensione

Visualizza su TensorFlow.org Esegui in Google Colab Visualizza l'origine su GitHubScarica quaderno

Impostare

!pip install -q tf_nightly
import tensorflow as tf
import numpy as np
from typing import Tuple, List, Mapping, Union, Optional
import tempfile

Tipi di estensione

I tipi definiti dall'utente possono rendere i progetti più leggibili, modulari e gestibili. Tuttavia, la maggior parte delle API TensorFlow ha un supporto molto limitato per i tipi Python definiti dall'utente. Ciò include API di alto livello (come Keras , tf.function , tf.SavedModel ) e API di livello inferiore (come tf.while_loop e tf.concat ). I tipi di estensione TensorFlow possono essere utilizzati per creare tipi orientati agli oggetti definiti dall'utente che funzionano perfettamente con le API di TensorFlow. Per creare un tipo di estensione, definisci semplicemente una classe Python con tf.experimental.ExtensionType come base e usa le annotazioni di tipo per specificare il tipo per ogni campo.

class TensorGraph(tf.experimental.ExtensionType):
  """A collection of labeled nodes connected by weighted edges."""
  edge_weights: tf.Tensor               # shape=[num_nodes, num_nodes]
  node_labels: Mapping[str, tf.Tensor]  # shape=[num_nodes]; dtype=any

class MaskedTensor(tf.experimental.ExtensionType):
  """A tensor paired with a boolean mask, indicating which values are valid."""
  values: tf.Tensor
  mask: tf.Tensor       # shape=values.shape; false for missing/invalid values.

class CSRSparseMatrix(tf.experimental.ExtensionType):
  """Compressed sparse row matrix (https://en.wikipedia.org/wiki/Sparse_matrix)."""
  values: tf.Tensor     # shape=[num_nonzero]; dtype=any
  col_index: tf.Tensor  # shape=[num_nonzero]; dtype=int64
  row_index: tf.Tensor  # shape=[num_rows+1]; dtype=int64

La classe base tf.experimental.ExtensionType funziona in modo simile a typing.NamedTuple e @dataclasses.dataclass della libreria Python standard. In particolare, aggiunge automaticamente un costruttore e metodi speciali (come __repr__ e __eq__ ) basati sulle annotazioni del tipo di campo.

In genere, i tipi di estensione tendono a rientrare in una di due categorie:

  • Strutture di dati , che raggruppano una raccolta di valori correlati e possono fornire operazioni utili basate su tali valori. Le strutture dei dati possono essere abbastanza generali (come l'esempio TensorGraph sopra); oppure possono essere altamente personalizzati per un modello specifico.

  • Tipi simili a tensori , che specializzano o estendono il concetto di "tensore". I tipi in questa categoria hanno un rank , una shape e solitamente un dtype ; e ha senso usarli con operazioni Tensor (come tf.stack , tf.add o tf.matmul ). MaskedTensor e CSRSparseMatrix sono esempi di tipi simili a tensori.

API supportate

I tipi di estensione sono supportati dalle seguenti API TensorFlow:

  • Keras : i tipi di estensione possono essere utilizzati come input e output per i Models e i Layers Keras.
  • tf.data.Dataset : i tipi di estensione possono essere inclusi nei Datasets di dati e restituiti dagli Iterators del set di dati.
  • Hub Tensorflow : i tipi di estensione possono essere utilizzati come ingressi e uscite per i moduli tf.hub .
  • SavedModel : i tipi di estensione possono essere utilizzati come input e output per le funzioni SavedModel .
  • tf.function : i tipi di estensione possono essere usati come argomenti e valori di ritorno per le funzioni racchiuse con il decoratore @tf.function .
  • while loops : i tipi di estensione possono essere usati come variabili di ciclo in tf.while_loop e possono essere usati come argomenti e valori di ritorno per il corpo del ciclo while.
  • conditionals : i tipi di estensione possono essere selezionati in modo condizionale utilizzando tf.cond e tf.case .
  • py_function : i tipi di estensione possono essere usati come argomenti e restituire valori per l'argomento func a tf.py_function .
  • Operazioni Tensor : i tipi di estensione possono essere estesi per supportare la maggior parte delle operazioni TensorFlow che accettano input Tensor (ad esempio, tf.matmul , tf.gather e tf.reduce_sum ). Per ulteriori informazioni, vedere la sezione " Invio " di seguito.
  • strategia di distribuzione : i tipi di estensione possono essere utilizzati come valori per replica.

Per maggiori dettagli, vedere la sezione "API TensorFlow che supportano ExtensionTypes" di seguito.

Requisiti

Tipi di campo

Tutti i campi (ovvero variabili di istanza) devono essere dichiarati e un'annotazione del tipo deve essere fornita per ogni campo. Sono supportate le seguenti annotazioni di tipo:

Tipo Esempio
Interi Python i: int
Python galleggia f: float
Stringhe Python s: str
Booleani Python b: bool
Python Nessuno n: None
Forme tensoriali shape: tf.TensorShape
Tipi di tensore dtype: tf.DType
Tensori t: tf.Tensor
Tipi di estensione mt: MyMaskedTensor
Tensori irregolari rt: tf.RaggedTensor
Tensori sparsi st: tf.SparseTensor
Fette indicizzate s: tf.IndexedSlices
Tensori opzionali o: tf.experimental.Optional
Digitare unioni int_or_float: typing.Union[int, float]
Tuple params: typing.Tuple[int, float, tf.Tensor, int]
Tuple di lunghezza variabile lengths: typing.Tuple[int, ...]
Mappature tags: typing.Mapping[str, tf.Tensor]
Valori facoltativi weight: typing.Optional[tf.Tensor]

Mutabilità

I tipi di estensione devono essere immutabili. Ciò garantisce che possano essere adeguatamente tracciati dai meccanismi di tracciamento dei grafici di TensorFlow. Se ti accorgi di voler mutare un valore del tipo di estensione, considera invece la definizione di metodi che trasformano i valori. Ad esempio, invece di definire un metodo set_mask per mutare un MaskedTensor , puoi definire un metodo replace_mask che restituisce un nuovo MaskedTensor :

class MaskedTensor(tf.experimental.ExtensionType):
  values: tf.Tensor
  mask: tf.Tensor

  def replace_mask(self, new_mask):
      self.values.shape.assert_is_compatible_with(new_mask.shape)
      return MaskedTensor(self.values, new_mask)

Funzionalità aggiunta da ExtensionType

La classe base ExtensionType fornisce le seguenti funzionalità:

  • Un costruttore ( __init__ ).
  • Un metodo di rappresentazione stampabile ( __repr__ ).
  • Operatori di uguaglianza e disuguaglianza ( __eq__ ).
  • Un metodo di convalida ( __validate__ ).
  • Immutabilità forzata.
  • Un TypeSpec nidificato.
  • Supporto per l'invio dell'API tensore.

Per ulteriori informazioni sulla personalizzazione di questa funzionalità, vedere la sezione "Personalizzazione dei tipi di estensione" di seguito.

Costruttore

Il costruttore aggiunto da ExtensionType accetta ogni campo come argomento denominato (nell'ordine in cui sono stati elencati nella definizione della classe). Questo costruttore verificherà il tipo di ogni parametro e lo convertirà ove necessario. In particolare, i campi Tensor vengono convertiti utilizzando tf.convert_to_tensor ; I campi delle Tuple vengono convertiti in tuple s; e i campi di Mapping vengono convertiti in dict immutabili.

class MaskedTensor(tf.experimental.ExtensionType):
  values: tf.Tensor
  mask: tf.Tensor

# Constructor takes one parameter for each field.
mt = MaskedTensor(values=[[1, 2, 3], [4, 5, 6]],
                  mask=[[True, True, False], [True, False, True]])

# Fields are type-checked and converted to the declared types.
# E.g., mt.values is converted to a Tensor.
print(mt.values)

tf.Tensor(
[[1 2 3]
 [4 5 6]], shape=(2, 3), dtype=int32)

Il costruttore genera un TypeError se un valore di campo non può essere convertito nel tipo dichiarato:

try:
  MaskedTensor([1, 2, 3], None)
except TypeError as e:
  print(f"Got expected TypeError: {e}")

Got expected TypeError: mask: expected a Tensor, got None

Il valore predefinito per un campo può essere specificato impostandone il valore a livello di classe:

class Pencil(tf.experimental.ExtensionType):
  color: str = "black"
  has_erasor: bool = True
  length: tf.Tensor = 1.0

Pencil()

Pencil(color='black', has_erasor=True, length=<tf.Tensor: shape=(), dtype=float32, numpy=1.0>)

Pencil(length=0.5, color="blue")

Pencil(color='blue', has_erasor=True, length=<tf.Tensor: shape=(), dtype=float32, numpy=0.5>)

Rappresentazione stampabile

ExtensionType aggiunge un metodo di rappresentazione stampabile predefinito ( __repr__ ) che include il nome della classe e il valore per ogni campo:

print(MaskedTensor(values=[1, 2, 3], mask=[True, True, False]))

MaskedTensor(values=<tf.Tensor: shape=(3,), dtype=int32, numpy=array([1, 2, 3], dtype=int32)>, mask=<tf.Tensor: shape=(3,), dtype=bool, numpy=array([ True,  True, False])>)

Operatori di uguaglianza

ExtensionType aggiunge operatori di uguaglianza predefiniti ( __eq__ e __ne__ ) che considerano due valori uguali se hanno lo stesso tipo e tutti i loro campi sono uguali. I campi tensoriale sono considerati uguali se hanno la stessa forma e sono uguali per tutti gli elementi.

a = MaskedTensor([1, 2], [True, False])
b = MaskedTensor([[3, 4], [5, 6]], [[False, True], [True, True]])
print(f"a == a: {a==a}")
print(f"a == b: {a==b}")
print(f"a == a.values: {a==a.values}")

a == a: True
a == b: False
a == a.values: False

Metodo di convalida

ExtensionType aggiunge un metodo __validate__ , che può essere sovrascritto per eseguire controlli di convalida sui campi. Viene eseguito dopo che il costruttore è stato chiamato e dopo che i campi sono stati controllati e convertiti nei loro tipi dichiarati, quindi può presumere che tutti i campi abbiano i loro tipi dichiarati.

L'esempio seguente aggiorna MaskedTensor per convalidare le shape s e dtype s dei suoi campi:

class MaskedTensor(tf.experimental.ExtensionType):
  """A tensor paired with a boolean mask, indicating which values are valid."""
  values: tf.Tensor
  mask: tf.Tensor
  def __validate__(self):
    self.values.shape.assert_is_compatible_with(self.mask.shape)
    assert self.mask.dtype.is_bool, 'mask.dtype must be bool'

try:
  MaskedTensor([1, 2, 3], [0, 1, 0])  # wrong dtype for mask.
except AssertionError as e:
  print(f"Got expected AssertionError: {e}")

Got expected AssertionError: mask.dtype must be bool

try:
  MaskedTensor([1, 2, 3], [True, False])  # shapes don't match.
except ValueError as e:
  print(f"Got expected ValueError: {e}")

Got expected ValueError: Shapes (3,) and (2,) are incompatible

Immutabilità forzata

ExtensionType sovrascrive i metodi __setattr__ e __delattr__ per prevenire la mutazione, assicurando che i valori del tipo di estensione siano immutabili.

mt = MaskedTensor([1, 2, 3], [True, False, True])

try:
  mt.mask = [True, True, True]
except AttributeError as e:
  print(f"Got expected AttributeError: {e}")

Got expected AttributeError: Cannot mutate attribute `mask` outside the custom constructor of ExtensionType.

try:
  mt.mask[0] = False
except TypeError as e:
  print(f"Got expected TypeError: {e}")

Got expected TypeError: 'tensorflow.python.framework.ops.EagerTensor' object does not support item assignment

try:
  del mt.mask
except AttributeError as e:
  print(f"Got expected AttributeError: {e}")

Got expected AttributeError: Cannot mutate attribute `mask` outside the custom constructor of ExtensionType.

Tipo nidificato Spec

Ogni classe ExtensionType ha una classe TypeSpec corrispondente, che viene creata automaticamente e archiviata come <extension_type_name>.Spec .

Questa classe acquisisce tutte le informazioni da un valore ad eccezione dei valori di eventuali tensori nidificati. In particolare, il TypeSpec per un valore viene creato sostituendo qualsiasi Tensor, ExtensionType o CompositeTensor nidificato con il relativo TypeSpec .

class Player(tf.experimental.ExtensionType):
  name: tf.Tensor
  attributes: Mapping[str, tf.Tensor]

anne = Player("Anne", {"height": 8.3, "speed": 28.1})
anne_spec = tf.type_spec_from_value(anne)
print(anne_spec.name)  # Records dtype and shape, but not the string value.
print(anne_spec.attributes)  # Records keys and TensorSpecs for values.

WARNING:tensorflow:Mapping types may not work well with tf.nest. Prefer using MutableMapping for <class 'tensorflow.python.framework.immutable_dict.ImmutableDict'>
TensorSpec(shape=(), dtype=tf.string, name=None)
ImmutableDict({'height': TensorSpec(shape=(), dtype=tf.float32, name=None), 'speed': TensorSpec(shape=(), dtype=tf.float32, name=None)})

I valori TypeSpec possono essere costruiti in modo esplicito oppure possono essere compilati da un valore ExtensionType usando tf.type_spec_from_value :

spec1 = Player.Spec(name=tf.TensorSpec([], tf.float32), attributes={})
spec2 = tf.type_spec_from_value(anne)

TypeSpec vengono utilizzati da TensorFlow per dividere i valori in un componente statico e un componente dinamico :

  • La componente statica (che è fissata al momento della costruzione del grafico) è codificata con un tf.TypeSpec .
  • La componente dinamica (che può variare ogni volta che viene eseguito il grafico) è codificata come un elenco di tf.Tensor s.

Ad esempio, tf.function la sua funzione avvolta ogni volta che un argomento ha un TypeSpec precedentemente non visto:

@tf.function
def anonymize_player(player):
  print("<<TRACING>>")
  return Player("<anonymous>", player.attributes)

# Function gets traced (first time the function has been called):
anonymize_player(Player("Anne", {"height": 8.3, "speed": 28.1}))

WARNING:tensorflow:Mapping types may not work well with tf.nest. Prefer using MutableMapping for <class 'tensorflow.python.framework.immutable_dict.ImmutableDict'>
WARNING:tensorflow:Mapping types may not work well with tf.nest. Prefer using MutableMapping for <class 'tensorflow.python.framework.immutable_dict.ImmutableDict'>
<<TRACING>>
Player(name=<tf.Tensor: shape=(), dtype=string, numpy=b'<anonymous>'>, attributes=ImmutableDict({'height': <tf.Tensor: shape=(), dtype=float32, numpy=8.3>, 'speed': <tf.Tensor: shape=(), dtype=float32, numpy=28.1>}))

# Function does NOT get traced (same TypeSpec: just tensor values changed)
anonymize_player(Player("Bart", {"height": 8.1, "speed": 25.3}))

Player(name=<tf.Tensor: shape=(), dtype=string, numpy=b'<anonymous>'>, attributes=ImmutableDict({'height': <tf.Tensor: shape=(), dtype=float32, numpy=8.1>, 'speed': <tf.Tensor: shape=(), dtype=float32, numpy=25.3>}))

# Function gets traced (new TypeSpec: keys for attributes changed):
anonymize_player(Player("Chuck", {"height": 11.0, "jump": 5.3}))

<<TRACING>>
Player(name=<tf.Tensor: shape=(), dtype=string, numpy=b'<anonymous>'>, attributes=ImmutableDict({'height': <tf.Tensor: shape=(), dtype=float32, numpy=11.0>, 'jump': <tf.Tensor: shape=(), dtype=float32, numpy=5.3>}))

Per ulteriori informazioni, vedere la Guida alla funzione tf .

Personalizzazione dei tipi di estensione

Oltre a dichiarare semplicemente i campi e i relativi tipi, i tipi di estensione possono:

  • Sostituisci la rappresentazione stampabile predefinita ( __repr__ ).
  • Definisci metodi.
  • Definire metodi di classe e metodi statici.
  • Definisci le proprietà.
  • Sostituisci il costruttore predefinito ( __init__ ).
  • Sostituisci l'operatore di uguaglianza predefinito ( __eq__ ).
  • Definire gli operatori (come __add__ e __lt__ ).
  • Dichiara i valori predefiniti per i campi.
  • Definisci le sottoclassi.

Sovrascrivere la rappresentazione stampabile predefinita

Puoi ignorare questo operatore di conversione di stringa predefinito per i tipi di estensione. L'esempio seguente aggiorna la classe MaskedTensor per generare una rappresentazione di stringa più leggibile quando i valori vengono stampati in modalità Eager.

class MaskedTensor(tf.experimental.ExtensionType):
  """A tensor paired with a boolean mask, indicating which values are valid."""
  values: tf.Tensor
  mask: tf.Tensor       # shape=values.shape; false for invalid values.

  def __repr__(self):
    return masked_tensor_str(self.values, self.mask)

def masked_tensor_str(values, mask):
  if isinstance(values, tf.Tensor):
    if hasattr(values, 'numpy') and hasattr(mask, 'numpy'):
      return f'<MaskedTensor {masked_tensor_str(values.numpy(), mask.numpy())}>'
    else:
      return f'MaskedTensor(values={values}, mask={mask})'
  if len(values.shape) == 1:
    items = [repr(v) if m else '_' for (v, m) in zip(values, mask)]
  else:
    items = [masked_tensor_str(v, m) for (v, m) in zip(values, mask)]
  return '[%s]' % ', '.join(items)

mt = MaskedTensor(values=[[1, 2, 3], [4, 5, 6]],
                  mask=[[True, True, False], [True, False, True]])
print(mt)

<MaskedTensor [[1, 2, _], [4, _, 6]]>

Definire i metodi

I tipi di estensione possono definire metodi, proprio come qualsiasi normale classe Python. Ad esempio, il tipo MaskedTensor potrebbe definire un metodo with_default che restituisce una copia di self con valori mascherati sostituiti da un determinato valore default . I metodi possono essere opzionalmente annotati con il decoratore @tf.function .

class MaskedTensor(tf.experimental.ExtensionType):
  values: tf.Tensor
  mask: tf.Tensor

  def with_default(self, default):
    return tf.where(self.mask, self.values, default)

MaskedTensor([1, 2, 3], [True, False, True]).with_default(0)

<tf.Tensor: shape=(3,), dtype=int32, numpy=array([1, 0, 3], dtype=int32)>

Definizione di metodi di classe e metodi statici

I tipi di estensione possono definire metodi usando i decoratori @classmethod e @staticmethod . Ad esempio, il tipo MaskedTensor potrebbe definire un metodo factory che maschera qualsiasi elemento con un determinato valore:

class MaskedTensor(tf.experimental.ExtensionType):
  values: tf.Tensor
  mask: tf.Tensor

  def __repr__(self):
    return masked_tensor_str(self.values, self.mask)

  @staticmethod
  def from_tensor_and_value_to_mask(values, value_to_mask):
    return MaskedTensor(values, values == value_to_mask)

x = tf.constant([[1, 0, 2], [3, 0, 0]])
MaskedTensor.from_tensor_and_value_to_mask(x, 0)

<MaskedTensor [[_, 0, _], [_, 0, 0]]>

Definizione delle proprietà

I tipi di estensione possono definire le proprietà usando il decoratore @property , proprio come qualsiasi normale classe Python. Ad esempio, il tipo MaskedTensor potrebbe definire una proprietà dtype che è una scorciatoia per il dtype dei valori:

class MaskedTensor(tf.experimental.ExtensionType):
  values: tf.Tensor
  mask: tf.Tensor

  @property
  def dtype(self):
    return self.values.dtype

MaskedTensor([1, 2, 3], [True, False, True]).dtype

tf.int32

Sovrascrivere il costruttore predefinito

Puoi sovrascrivere il costruttore predefinito per i tipi di estensione. I costruttori personalizzati devono impostare un valore per ogni campo dichiarato; e dopo la restituzione del costruttore personalizzato, tutti i campi verranno controllati e i valori verranno convertiti come descritto sopra.

class Toy(tf.experimental.ExtensionType):
  name: str
  price: tf.Tensor
  def __init__(self, name, price, discount=0):
    self.name = name
    self.price = price * (1 - discount)

print(Toy("ball", 5.0, discount=0.2))  # On sale -- 20% off!

Toy(name='ball', price=<tf.Tensor: shape=(), dtype=float32, numpy=4.0>)

In alternativa, potresti considerare di lasciare il costruttore predefinito così com'è, ma di aggiungere uno o più metodi di fabbrica. Per esempio:

class Toy(tf.experimental.ExtensionType):
  name: str
  price: tf.Tensor

  @staticmethod
  def new_toy_with_discount(name, price, discount):
    return Toy(name, price * (1 - discount))

print(Toy.new_toy_with_discount("ball", 5.0, discount=0.2))

Toy(name='ball', price=<tf.Tensor: shape=(), dtype=float32, numpy=4.0>)

Sovrascrivere l'operatore di uguaglianza predefinito ( __eq__ )

È possibile sostituire l'operatore __eq__ predefinito per i tipi di estensione. L'esempio seguente aggiorna MaskedTensor per ignorare gli elementi mascherati durante il confronto per l'uguaglianza.

class MaskedTensor(tf.experimental.ExtensionType):
  values: tf.Tensor
  mask: tf.Tensor

  def __repr__(self):
    return masked_tensor_str(self.values, self.mask)

  def __eq__(self, other):
    result = tf.math.equal(self.values, other.values)
    result = result | ~(self.mask & other.mask)
    return tf.reduce_all(result)

x = MaskedTensor([1, 2, 3, 4], [True, True, False, True])
y = MaskedTensor([5, 2, 0, 4], [False, True, False, True])
print(x == y)

tf.Tensor(True, shape=(), dtype=bool)

Utilizzo di riferimenti diretti

Se il tipo per un campo non è stato ancora definito, puoi invece utilizzare una stringa contenente il nome del tipo. Nell'esempio seguente, la stringa "Node" viene utilizzata per annotare il campo children poiché il tipo di Node non è stato ancora (completamente) definito.

class Node(tf.experimental.ExtensionType):
  value: tf.Tensor
  children: Tuple["Node", ...] = ()

Node(3, [Node(5), Node(2)])

Node(value=<tf.Tensor: shape=(), dtype=int32, numpy=3>, children=(Node(value=<tf.Tensor: shape=(), dtype=int32, numpy=5>, children=()), Node(value=<tf.Tensor: shape=(), dtype=int32, numpy=2>, children=())))

Definizione delle sottoclassi

I tipi di estensione possono essere sottoclassi usando la sintassi standard di Python. Le sottoclassi del tipo di estensione possono aggiungere nuovi campi, metodi e proprietà; e può sovrascrivere il costruttore, la rappresentazione stampabile e l'operatore di uguaglianza. L'esempio seguente definisce una classe TensorGraph di base che utilizza tre campi Tensor per codificare un insieme di bordi tra i nodi. Quindi definisce una sottoclasse che aggiunge un campo Tensor per registrare un "valore della caratteristica" per ogni nodo. La sottoclasse definisce anche un metodo per propagare i valori delle caratteristiche lungo i bordi.

class TensorGraph(tf.experimental.ExtensionType):
  num_nodes: tf.Tensor
  edge_src: tf.Tensor   # edge_src[e] = index of src node for edge e.
  edge_dst: tf.Tensor   # edge_dst[e] = index of dst node for edge e.

class TensorGraphWithNodeFeature(TensorGraph):
  node_features: tf.Tensor  # node_features[n] = feature value for node n.

  def propagate_features(self, weight=1.0) -> 'TensorGraphWithNodeFeature':
    updates = tf.gather(self.node_features, self.edge_src) * weight
    new_node_features = tf.tensor_scatter_nd_add(
        self.node_features, tf.expand_dims(self.edge_dst, 1), updates)
    return TensorGraphWithNodeFeature(
        self.num_nodes, self.edge_src, self.edge_dst, new_node_features)

g = TensorGraphWithNodeFeature(  # Edges: 0->1, 4->3, 2->2, 2->1
    num_nodes=5, edge_src=[0, 4, 2, 2], edge_dst=[1, 3, 2, 1],
    node_features=[10.0, 0.0, 2.0, 5.0, -1.0, 0.0])

print("Original features:", g.node_features)
print("After propagating:", g.propagate_features().node_features)

Original features: tf.Tensor([10.  0.  2.  5. -1.  0.], shape=(6,), dtype=float32)
After propagating: tf.Tensor([10. 12.  4.  4. -1.  0.], shape=(6,), dtype=float32)

Definizione dei campi privati

I campi di un tipo di estensione possono essere contrassegnati come privati ​​anteponendoli a un trattino basso (seguendo le convenzioni standard di Python). Ciò non influisce in alcun modo sul modo in cui TensorFlow tratta i campi; ma serve semplicemente come segnale a tutti gli utenti del tipo di estensione che quei campi sono privati.

Personalizzazione di TypeSpec di ExtensionType

Ogni classe ExtensionType ha una classe TypeSpec corrispondente, che viene creata automaticamente e archiviata come <extension_type_name>.Spec . Per ulteriori informazioni, vedere la sezione "TipoSpec nidificato" sopra.

Per personalizzare TypeSpec , definire semplicemente la propria classe nidificata denominata Spec e ExtensionType la utilizzerà come base per la TypeSpec costruita automaticamente. Puoi personalizzare la classe Spec :

  • Sovrascrivere la rappresentazione stampabile predefinita.
  • Sovrascrivere il costruttore predefinito.
  • Definizione di metodi, metodi di classe, metodi statici e proprietà.

L'esempio seguente personalizza la classe MaskedTensor.Spec per semplificarne l'utilizzo:

class MaskedTensor(tf.experimental.ExtensionType):
  values: tf.Tensor
  mask: tf.Tensor

  shape = property(lambda self: self.values.shape)
  dtype = property(lambda self: self.values.dtype)

  def __repr__(self):
    return masked_tensor_str(self.values, self.mask)

  def with_values(self, new_values):
    return MaskedTensor(new_values, self.mask)

  class Spec:
    def __init__(self, shape, dtype=tf.float32):
      self.values = tf.TensorSpec(shape, dtype)
      self.mask = tf.TensorSpec(shape, tf.bool)

    def __repr__(self):
      return f"MaskedTensor.Spec(shape={self.shape}, dtype={self.dtype})"

    shape = property(lambda self: self.values.shape)
    dtype = property(lambda self: self.values.dtype)

Invio dell'API tensore

I tipi di estensione possono essere "tensoriali", nel senso che specializzano o estendono l'interfaccia definita dal tipo tf.Tensor . Esempi di tipi di estensioni simili a tensori includono RaggedTensor , SparseTensor e MaskedTensor . I decoratori di spedizione possono essere utilizzati per sovrascrivere il comportamento predefinito delle operazioni TensorFlow quando applicati a tipi di estensione simili a tensori. TensorFlow definisce attualmente tre decoratori di spedizione:

Invio per una singola API

Il decoratore tf.experimental.dispatch_for_api sovrascrive il comportamento predefinito di un'operazione TensorFlow specificata quando viene chiamata con la firma specificata. Ad esempio, puoi utilizzare questo decoratore per specificare come tf.stack dovrebbe elaborare i valori MaskedTensor :

@tf.experimental.dispatch_for_api(tf.stack)
def masked_stack(values: List[MaskedTensor], axis = 0):
  return MaskedTensor(tf.stack([v.values for v in values], axis),
                      tf.stack([v.mask for v in values], axis))

Questo sovrascrive l'implementazione predefinita per tf.stack ogni volta che viene chiamato con un elenco di valori MaskedTensor (poiché l'argomento values è annotato con typing.List[MaskedTensor] ):

x = MaskedTensor([1, 2, 3], [True, True, False])
y = MaskedTensor([4, 5, 6], [False, True, True])
tf.stack([x, y])

<MaskedTensor [[1, 2, _], [_, 5, 6]]>

Per consentire a tf.stack di gestire elenchi di valori misti MaskedTensor e Tensor , puoi perfezionare l'annotazione del tipo per il parametro values e aggiornare il corpo della funzione in modo appropriato:

tf.experimental.unregister_dispatch_for(masked_stack)

def convert_to_masked_tensor(x):
  if isinstance(x, MaskedTensor):
    return x
  else:
    return MaskedTensor(x, tf.ones_like(x, tf.bool))

@tf.experimental.dispatch_for_api(tf.stack)
def masked_stack_v2(values: List[Union[MaskedTensor, tf.Tensor]], axis = 0):
  values = [convert_to_masked_tensor(v) for v in values]
  return MaskedTensor(tf.stack([v.values for v in values], axis),
                      tf.stack([v.mask for v in values], axis))
x = MaskedTensor([1, 2, 3], [True, True, False])
y = tf.constant([4, 5, 6])
tf.stack([x, y, x])

<MaskedTensor [[1, 2, _], [4, 5, 6], [1, 2, _]]>

Per un elenco di API che possono essere sovrascritte, consulta la documentazione API per tf.experimental.dispatch_for_api .

Invio per tutte le API elementwise unarie

Il decoratore tf.experimental.dispatch_for_unary_elementwise_apis sovrascrive il comportamento predefinito di tutte le operazioni elementari unarie (come tf.math.cos ) ogni volta che il valore del primo argomento (in genere chiamato x ) corrisponde all'annotazione del tipo x_type . La funzione decorata dovrebbe avere due argomenti:

  • api_func : una funzione che accetta un singolo parametro ed esegue l'operazione a livello di elemento (ad esempio, tf.abs ).
  • x : il primo argomento dell'operazione elementwise.

L'esempio seguente aggiorna tutte le operazioni unarie sugli elementi per gestire il tipo MaskedTensor :

@tf.experimental.dispatch_for_unary_elementwise_apis(MaskedTensor)
 def masked_tensor_unary_elementwise_api_handler(api_func, x):
   return MaskedTensor(api_func(x.values), x.mask)

Questa funzione verrà ora utilizzata ogni volta che viene chiamata un'operazione unaria sugli elementi su un MaskedTensor .

x = MaskedTensor([1, -2, -3], [True, False, True])
 print(tf.abs(x))

<MaskedTensor [1, _, 3]>

print(tf.ones_like(x, dtype=tf.float32))

<MaskedTensor [1.0, _, 1.0]>

Invio per tutte le API elementwise binarie

Allo stesso modo, tf.experimental.dispatch_for_binary_elementwise_apis può essere utilizzato per aggiornare tutte le operazioni elementwise binarie per gestire il tipo MaskedTensor :

@tf.experimental.dispatch_for_binary_elementwise_apis(MaskedTensor, MaskedTensor)
def masked_tensor_binary_elementwise_api_handler(api_func, x, y):
  return MaskedTensor(api_func(x.values, y.values), x.mask & y.mask)

x = MaskedTensor([1, -2, -3], [True, False, True])
y = MaskedTensor([[4], [5]], [[True], [False]])
tf.math.add(x, y)

<MaskedTensor [[5, _, 1], [_, _, _]]>

Per un elenco delle API elementwise che sono state sovrascritte, consulta la documentazione API per tf.experimental.dispatch_for_unary_elementwise_apis e tf.experimental.dispatch_for_binary_elementwise_apis .

Tipi di estensione batch

Un ExtensionType è inviabile in batch se è possibile utilizzare una singola istanza per rappresentare un batch di valori. In genere, ciò si ottiene aggiungendo dimensioni batch a tutti i Tensor nidificati. Le seguenti API TensorFlow richiedono che qualsiasi input di tipo di estensione sia batch:

Per impostazione predefinita, BatchableExtensionType crea valori in batch inviando in batch qualsiasi Tensor s, CompositeTensor s ed ExtensionType s nidificato. Se questo non è appropriato per la tua classe, dovrai usare tf.experimental.ExtensionTypeBatchEncoder per sovrascrivere questo comportamento predefinito. Ad esempio, non sarebbe appropriato creare un batch di valori tf.SparseTensor semplicemente impilando i singoli campi sparse tensors' values , indices e dense_shape -- nella maggior parte dei casi, non è possibile impilare questi tensori, poiché hanno forme incompatibili ; e anche se potessi, il risultato non sarebbe uno SparseTensor valido.

Esempio BatchableExtensionType: Rete

Ad esempio, si consideri una semplice classe Network utilizzata per il bilanciamento del carico, che tiene traccia di quanto lavoro è rimasto da fare su ciascun nodo e quanta larghezza di banda è disponibile per spostare il lavoro tra i nodi:

class Network(tf.experimental.ExtensionType):  # This version is not batchable.
  work: tf.Tensor       # work[n] = work left to do at node n
  bandwidth: tf.Tensor  # bandwidth[n1, n2] = bandwidth from n1->n2

net1 = Network([5., 3, 8], [[0., 2, 0], [2, 0, 3], [0, 3, 0]])
net2 = Network([3., 4, 2], [[0., 2, 2], [2, 0, 2], [2, 2, 0]])

Per rendere questo tipo batchable, modificare il tipo di base in BatchableExtensionType e regolare la forma di ogni campo per includere dimensioni batch facoltative. L'esempio seguente aggiunge anche un campo shape per tenere traccia della forma batch. Questo campo shape non è richiesto da tf.data.Dataset o tf.map_fn , ma è richiesto da tf.Keras .

class Network(tf.experimental.BatchableExtensionType):
  shape: tf.TensorShape  # batch shape.  A single network has shape=[].
  work: tf.Tensor        # work[*shape, n] = work left to do at node n
  bandwidth: tf.Tensor   # bandwidth[*shape, n1, n2] = bandwidth from n1->n2

  def __init__(self, work, bandwidth):
    self.work = tf.convert_to_tensor(work)
    self.bandwidth = tf.convert_to_tensor(bandwidth)
    work_batch_shape = self.work.shape[:-1]
    bandwidth_batch_shape = self.bandwidth.shape[:-2]
    self.shape = work_batch_shape.merge_with(bandwidth_batch_shape)

  def __repr__(self):
    return network_repr(self)

def network_repr(network):
  work = network.work
  bandwidth = network.bandwidth
  if hasattr(work, 'numpy'):
    work = ' '.join(str(work.numpy()).split())
  if hasattr(bandwidth, 'numpy'):
    bandwidth = ' '.join(str(bandwidth.numpy()).split())
  return (f"<Network shape={network.shape} work={work} bandwidth={bandwidth}>")

net1 = Network([5., 3, 8], [[0., 2, 0], [2, 0, 3], [0, 3, 0]])
net2 = Network([3., 4, 2], [[0., 2, 2], [2, 0, 2], [2, 2, 0]])
batch_of_networks = Network(
    work=tf.stack([net1.work, net2.work]),
    bandwidth=tf.stack([net1.bandwidth, net2.bandwidth]))
print(f"net1={net1}")
print(f"net2={net2}")
print(f"batch={batch_of_networks}")

net1=<Network shape=() work=[5. 3. 8.] bandwidth=[[0. 2. 0.] [2. 0. 3.] [0. 3. 0.]]>
net2=<Network shape=() work=[3. 4. 2.] bandwidth=[[0. 2. 2.] [2. 0. 2.] [2. 2. 0.]]>
batch=<Network shape=(2,) work=[[5. 3. 8.] [3. 4. 2.]] bandwidth=[[[0. 2. 0.] [2. 0. 3.] [0. 3. 0.]] [[0. 2. 2.] [2. 0. 2.] [2. 2. 0.]]]>

È quindi possibile utilizzare tf.data.Dataset per scorrere un batch di reti:

dataset = tf.data.Dataset.from_tensor_slices(batch_of_networks)
for i, network in enumerate(dataset):
  print(f"Batch element {i}: {network}")

Batch element 0: <Network shape=() work=[5. 3. 8.] bandwidth=[[0. 2. 0.] [2. 0. 3.] [0. 3. 0.]]>
Batch element 1: <Network shape=() work=[3. 4. 2.] bandwidth=[[0. 2. 2.] [2. 0. 2.] [2. 2. 0.]]>

E puoi anche usare map_fn per applicare una funzione a ciascun elemento batch:

def balance_work_greedy(network):
  delta = (tf.expand_dims(network.work, -1) - tf.expand_dims(network.work, -2))
  delta /= 4
  delta = tf.maximum(tf.minimum(delta, network.bandwidth), -network.bandwidth)
  new_work = network.work + tf.reduce_sum(delta, -1)
  return Network(new_work, network.bandwidth)

tf.map_fn(balance_work_greedy, batch_of_networks)

<Network shape=(2,) work=[[5.5 1.25 9.25] [3. 4.75 1.25]] bandwidth=[[[0. 2. 0.] [2. 0. 3.] [0. 3. 0.]] [[0. 2. 2.] [2. 0. 2.] [2. 2. 0.]]]>

API TensorFlow che supportano ExtensionTypes

@tf.funzione

tf.function è un decoratore che precalcola i grafici TensorFlow per le funzioni Python, che possono migliorare sostanzialmente le prestazioni del codice TensorFlow. I valori del tipo di estensione possono essere utilizzati in modo trasparente con @tf.function -decorated functions.

class Pastry(tf.experimental.ExtensionType):
  sweetness: tf.Tensor  # 2d embedding that encodes sweetness
  chewiness: tf.Tensor  # 2d embedding that encodes chewiness

@tf.function
def combine_pastry_features(x: Pastry):
  return (x.sweetness + x.chewiness) / 2

cookie = Pastry(sweetness=[1.2, 0.4], chewiness=[0.8, 0.2])
combine_pastry_features(cookie)

<tf.Tensor: shape=(2,), dtype=float32, numpy=array([1. , 0.3], dtype=float32)>

Se desideri specificare in modo esplicito input_signature per tf.function , puoi farlo utilizzando TypeSpec del tipo di estensione.

pastry_spec = Pastry.Spec(tf.TensorSpec([2]), tf.TensorSpec(2))

@tf.function(input_signature=[pastry_spec])
def increase_sweetness(x: Pastry, delta=1.0):
  return Pastry(x.sweetness + delta, x.chewiness)

increase_sweetness(cookie)

Pastry(sweetness=<tf.Tensor: shape=(2,), dtype=float32, numpy=array([2.2, 1.4], dtype=float32)>, chewiness=<tf.Tensor: shape=(2,), dtype=float32, numpy=array([0.8, 0.2], dtype=float32)>)

Funzioni concrete

Le funzioni concrete incapsulano i singoli grafici tracciati che sono costruiti da tf.function . I tipi di estensione possono essere utilizzati in modo trasparente con funzioni concrete.

cf = combine_pastry_features.get_concrete_function(pastry_spec)
cf(cookie)

<tf.Tensor: shape=(2,), dtype=float32, numpy=array([1. , 0.3], dtype=float32)>

Controllare le operazioni di flusso

I tipi di estensione sono supportati dalle operazioni del flusso di controllo di TensorFlow:

# Example: using tf.cond to select between two MaskedTensors.  Note that the
# two MaskedTensors don't need to have the same shape.
a = MaskedTensor([1., 2, 3], [True, False, True])
b = MaskedTensor([22., 33, 108, 55], [True, True, True, False])
condition = tf.constant(True)
print(tf.cond(condition, lambda: a, lambda: b))

<MaskedTensor [1.0, _, 3.0]>

# Example: using tf.while_loop with MaskedTensor.
cond = lambda i, _: i < 10
def body(i, mt):
  return i + 1, mt.with_values(mt.values + 3 / 7)
print(tf.while_loop(cond, body, [0, b])[1])

<MaskedTensor [26.285717, 37.285698, 112.285736, _]>

Flusso di controllo autografo

I tipi di estensione sono supportati anche dalle istruzioni del flusso di controllo in tf.function (usando l'autografo). Nell'esempio seguente, le istruzioni if e for vengono automaticamente convertite in operazioni tf.cond e tf.while_loop , che supportano i tipi di estensione.

@tf.function
def fn(x, b):
  if b:
    x = MaskedTensor(x, tf.less(x, 0))
  else:
    x = MaskedTensor(x, tf.greater(x, 0))
  for i in tf.range(5 if b else 7):
    x = x.with_values(x.values + 1 / 2)
  return x

print(fn(tf.constant([1., -2, 3]), tf.constant(True)))
print(fn(tf.constant([1., -2, 3]), tf.constant(False)))

<MaskedTensor [_, 0.5, _]>
<MaskedTensor [4.5, _, 6.5]>

Cheras

tf.keras è l'API di alto livello di TensorFlow per la creazione e il training di modelli di deep learning. I tipi di estensione possono essere passati come input a un modello Keras, passati tra livelli Keras e restituiti dai modelli Keras. Keras attualmente pone due requisiti sui tipi di estensione:

  • Devono essere batch (vedi "Tipi di estensione in batch" sopra).
  • Deve avere un campo o una proprietà denominata shape . shape[0] è considerata la dimensione batch.

Le due sottosezioni seguenti forniscono esempi che mostrano come utilizzare i tipi di estensione con Keras.

Esempio Keras: Network

Per il primo esempio, considera la classe Network definita nella sezione "Batchable ExtensionTypes" precedente, che può essere utilizzata per il lavoro di bilanciamento del carico tra i nodi. La sua definizione è qui ripetuta:

class Network(tf.experimental.BatchableExtensionType):
  shape: tf.TensorShape  # batch shape.  A single network has shape=[].
  work: tf.Tensor        # work[*shape, n] = work left to do at node n
  bandwidth: tf.Tensor   # bandwidth[*shape, n1, n2] = bandwidth from n1->n2

  def __init__(self, work, bandwidth):
    self.work = tf.convert_to_tensor(work)
    self.bandwidth = tf.convert_to_tensor(bandwidth)
    work_batch_shape = self.work.shape[:-1]
    bandwidth_batch_shape = self.bandwidth.shape[:-2]
    self.shape = work_batch_shape.merge_with(bandwidth_batch_shape)

  def __repr__(self):
    return network_repr(self)

single_network = Network(  # A single network w/ 4 nodes.
    work=[8.0, 5, 12, 2],
    bandwidth=[[0.0, 1, 2, 2], [1, 0, 0, 2], [2, 0, 0, 1], [2, 2, 1, 0]])

batch_of_networks = Network(  # Batch of 2 networks, each w/ 2 nodes.
    work=[[8.0, 5], [3, 2]],
    bandwidth=[[[0.0, 1], [1, 0]], [[0, 2], [2, 0]]])

È possibile definire un nuovo livello Keras che elabora i messaggi di Network .

class BalanceNetworkLayer(tf.keras.layers.Layer):
  """Layer that balances work between nodes in a network.

  Shifts work from more busy nodes to less busy nodes, constrained by bandwidth.
  """
  def call(self, inputs):
    # This function is defined above, in "Batchable ExtensionTypes" section.
    return balance_work_greedy(inputs)

È quindi possibile utilizzare questi livelli per creare un modello semplice. Per inserire un ExtensionType in un modello, puoi utilizzare un livello tf.keras.layer.Input con type_spec impostato su TypeSpec del tipo di estensione. Se il modello Keras verrà utilizzato per elaborare i batch, type_spec deve includere la dimensione batch.

input_spec = Network.Spec(shape=None,
                          work=tf.TensorSpec(None, tf.float32),
                          bandwidth=tf.TensorSpec(None, tf.float32))
model = tf.keras.Sequential([
    tf.keras.layers.Input(type_spec=input_spec),
    BalanceNetworkLayer(),
    ])

Infine, puoi applicare il modello a una singola rete ea un batch di reti.

model(single_network)

<Network shape=() work=[ 9.25 5. 14. -1.25] bandwidth=[[0. 1. 2. 2.] [1. 0. 0. 2.] [2. 0. 0. 1.] [2. 2. 1. 0.]]>

model(batch_of_networks)

<Network shape=(2,) work=[[8.75 4.25] [3.25 1.75]] bandwidth=[[[0. 1.] [1. 0.]] [[0. 2.] [2. 0.]]]>

Esempio Keras: MaskedTensor

In questo esempio, MaskedTensor viene esteso per supportare Keras . shape è definito come una proprietà che viene calcolata dal campo dei values . Keras richiede che tu aggiunga questa proprietà sia al tipo di estensione che alla sua TypeSpec . MaskedTensor definisce anche una variabile __name__ , che sarà richiesta per la serializzazione di SavedModel (sotto).

class MaskedTensor(tf.experimental.BatchableExtensionType):
  # __name__ is required for serialization in SavedModel; see below for details.
  __name__ = 'extension_type_colab.MaskedTensor'

  values: tf.Tensor
  mask: tf.Tensor

  shape = property(lambda self: self.values.shape)
  dtype = property(lambda self: self.values.dtype)

  def with_default(self, default):
    return tf.where(self.mask, self.values, default)

  def __repr__(self):
    return masked_tensor_str(self.values, self.mask)

  class Spec:
    def __init__(self, shape, dtype=tf.float32):
      self.values = tf.TensorSpec(shape, dtype)
      self.mask = tf.TensorSpec(shape, tf.bool)

    shape = property(lambda self: self.values.shape)
    dtype = property(lambda self: self.values.dtype)

    def with_shape(self):
      return MaskedTensor.Spec(tf.TensorSpec(shape, self.values.dtype),
                               tf.TensorSpec(shape, self.mask.dtype))

Successivamente, i decoratori di spedizione vengono utilizzati per sovrascrivere il comportamento predefinito di diverse API TensorFlow. Poiché queste API sono utilizzate da livelli Keras standard (come il livello Dense ), l'override di questi ci consentirà di utilizzare quei livelli con MaskedTensor . Ai fini di questo esempio, matmul per tensori mascherati è definito per trattare i valori mascherati come zeri (cioè, per non includerli nel prodotto).

@tf.experimental.dispatch_for_unary_elementwise_apis(MaskedTensor)
def unary_elementwise_op_handler(op, x):
 return MaskedTensor(op(x.values), x.mask)

@tf.experimental.dispatch_for_binary_elementwise_apis(
    Union[MaskedTensor, tf.Tensor],
    Union[MaskedTensor, tf.Tensor])
def binary_elementwise_op_handler(op, x, y):
  x = convert_to_masked_tensor(x)
  y = convert_to_masked_tensor(y)
  return MaskedTensor(op(x.values, y.values), x.mask & y.mask)

@tf.experimental.dispatch_for_api(tf.matmul)
def masked_matmul(a: MaskedTensor, b,
                  transpose_a=False, transpose_b=False,
                  adjoint_a=False, adjoint_b=False,
                  a_is_sparse=False, b_is_sparse=False,
                  output_type=None):
  if isinstance(a, MaskedTensor):
    a = a.with_default(0)
  if isinstance(b, MaskedTensor):
    b = b.with_default(0)
  return tf.matmul(a, b, transpose_a, transpose_b, adjoint_a,
                   adjoint_b, a_is_sparse, b_is_sparse, output_type)

È quindi possibile costruire un modello Keras che accetta input MaskedTensor , utilizzando i livelli Keras standard:

input_spec = MaskedTensor.Spec([None, 2], tf.float32)

masked_tensor_model = tf.keras.Sequential([
    tf.keras.layers.Input(type_spec=input_spec),
    tf.keras.layers.Dense(16, activation="relu"),
    tf.keras.layers.Dense(1)])
masked_tensor_model.compile(loss='binary_crossentropy', optimizer='rmsprop')

a = MaskedTensor([[1., 2], [3, 4], [5, 6]],
                  [[True, False], [False, True], [True, True]])
masked_tensor_model.fit(a, tf.constant([[1], [0], [1]]), epochs=3)
print(masked_tensor_model(a))

Epoch 1/3
1/1 [==============================] - 1s 955ms/step - loss: 10.2833
Epoch 2/3
1/1 [==============================] - 0s 5ms/step - loss: 10.2833
Epoch 3/3
1/1 [==============================] - 0s 5ms/step - loss: 10.2833
tf.Tensor(
[[-0.09944128]
 [-0.7225147 ]
 [-1.3020657 ]], shape=(3, 1), dtype=float32)

Modello salvato

Un SavedModel è un programma TensorFlow serializzato, che include sia i pesi che il calcolo. Può essere costruito da un modello Keras o da un modello personalizzato. In entrambi i casi, i tipi di estensione possono essere utilizzati in modo trasparente con le funzioni ei metodi definiti da un SavedModel.

SavedModel può salvare modelli, livelli e funzioni che elaborano i tipi di estensione, purché i tipi di estensione abbiano un campo __name__ . Questo nome viene utilizzato per registrare il tipo di estensione, quindi può essere individuato quando il modello viene caricato.

Esempio: salvataggio di un modello Keras

I modelli Keras che utilizzano i tipi di estensione possono essere salvati utilizzando SavedModel .

masked_tensor_model_path = tempfile.mkdtemp()
tf.saved_model.save(masked_tensor_model, masked_tensor_model_path)
imported_model = tf.saved_model.load(masked_tensor_model_path)
imported_model(a)

2021-11-06 01:25:14.285250: W tensorflow/python/util/util.cc:368] Sets are not currently considered sequences, but this may change in the future, so consider avoiding using them.
WARNING:absl:Function `_wrapped_model` contains input name(s) args_0 with unsupported characters which will be renamed to args_0_1 in the SavedModel.
INFO:tensorflow:Assets written to: /tmp/tmp3ceuupv9/assets
INFO:tensorflow:Assets written to: /tmp/tmp3ceuupv9/assets
<tf.Tensor: shape=(3, 1), dtype=float32, numpy=
array([[-0.09944128],
       [-0.7225147 ],
       [-1.3020657 ]], dtype=float32)>

Esempio: salvataggio di un modello personalizzato

SavedModel può essere utilizzato anche per salvare sottoclassi tf.Module personalizzate con funzioni che elaborano i tipi di estensione.

class CustomModule(tf.Module):
  def __init__(self, variable_value):
    super().__init__()
    self.v = tf.Variable(variable_value)

  @tf.function
  def grow(self, x: MaskedTensor):
    """Increase values in `x` by multiplying them by `self.v`."""
    return MaskedTensor(x.values * self.v, x.mask)

module = CustomModule(100.0)

module.grow.get_concrete_function(MaskedTensor.Spec(shape=None,
                                                    dtype=tf.float32))
custom_module_path = tempfile.mkdtemp()
tf.saved_model.save(module, custom_module_path)
imported_model = tf.saved_model.load(custom_module_path)
imported_model.grow(MaskedTensor([1., 2, 3], [False, True, False]))

INFO:tensorflow:Assets written to: /tmp/tmp2x8zq5kb/assets
INFO:tensorflow:Assets written to: /tmp/tmp2x8zq5kb/assets
<MaskedTensor [_, 200.0, _]>

Caricamento di un SavedModel quando ExtensionType non è disponibile

Se carichi un SavedModel che utilizza ExtensionType , ma tale ExtensionType non è disponibile (cioè non è stato importato), vedrai un avviso e TensorFlow tornerà all'utilizzo di un oggetto "tipo di estensione anonimo". Questo oggetto avrà gli stessi campi del tipo originale, ma mancherà di ulteriori personalizzazioni aggiunte per il tipo, come metodi o proprietà personalizzate.

Utilizzo di ExtensionTypes con il servizio TensorFlow

Attualmente, il servizio TensorFlow (e altri consumer del dizionario "firme" di SavedModel) richiedono che tutti gli input e gli output siano tensori grezzi. Se desideri utilizzare il servizio TensorFlow con un modello che utilizza tipi di estensione, puoi aggiungere metodi wrapper che compongono o scompongono i valori del tipo di estensione dai tensori. Per esempio:

class CustomModuleWrapper(tf.Module):
  def __init__(self, variable_value):
    super().__init__()
    self.v = tf.Variable(variable_value)

  @tf.function
  def var_weighted_mean(self, x: MaskedTensor):
    """Mean value of unmasked values in x, weighted by self.v."""
    x = MaskedTensor(x.values * self.v, x.mask)
    return (tf.reduce_sum(x.with_default(0)) /
            tf.reduce_sum(tf.cast(x.mask, x.dtype)))

  @tf.function()
  def var_weighted_mean_wrapper(self, x_values, x_mask):
    """Raw tensor wrapper for var_weighted_mean."""
    return self.var_weighted_mean(MaskedTensor(x_values, x_mask))

module = CustomModuleWrapper([3., 2., 8., 5.])

module.var_weighted_mean_wrapper.get_concrete_function(
    tf.TensorSpec(None, tf.float32), tf.TensorSpec(None, tf.bool))
custom_module_path = tempfile.mkdtemp()
tf.saved_model.save(module, custom_module_path)
imported_model = tf.saved_model.load(custom_module_path)
x = MaskedTensor([1., 2., 3., 4.], [False, True, False, True])
imported_model.var_weighted_mean_wrapper(x.values, x.mask)

INFO:tensorflow:Assets written to: /tmp/tmpxhh4zh0i/assets
INFO:tensorflow:Assets written to: /tmp/tmpxhh4zh0i/assets
<tf.Tensor: shape=(), dtype=float32, numpy=12.0>

Set di dati

tf.data è un'API che consente di creare pipeline di input complesse da parti semplici e riutilizzabili. La sua struttura di dati di base è tf.data.Dataset , che rappresenta una sequenza di elementi, in cui ogni elemento è costituito da uno o più componenti.

Creazione di set di dati con tipi di estensione

I set di dati possono essere creati dai valori del tipo di estensione utilizzando Dataset.from_tensors , Dataset.from_tensor_slices o Dataset.from_generator :

ds = tf.data.Dataset.from_tensors(Pastry(5, 5))
iter(ds).next()

Pastry(sweetness=<tf.Tensor: shape=(), dtype=int32, numpy=5>, chewiness=<tf.Tensor: shape=(), dtype=int32, numpy=5>)

mt = MaskedTensor(tf.reshape(range(20), [5, 4]), tf.ones([5, 4]))
ds = tf.data.Dataset.from_tensor_slices(mt)
for value in ds:
  print(value)

<MaskedTensor [0, 1, 2, 3]>
<MaskedTensor [4, 5, 6, 7]>
<MaskedTensor [8, 9, 10, 11]>
<MaskedTensor [12, 13, 14, 15]>
<MaskedTensor [16, 17, 18, 19]>

def value_gen():
  for i in range(2, 7):
    yield MaskedTensor(range(10), [j%i != 0 for j in range(10)])

ds = tf.data.Dataset.from_generator(
    value_gen, output_signature=MaskedTensor.Spec(shape=[10], dtype=tf.int32))
for value in ds:
  print(value)

<MaskedTensor [_, 1, _, 3, _, 5, _, 7, _, 9]>
<MaskedTensor [_, 1, 2, _, 4, 5, _, 7, 8, _]>
<MaskedTensor [_, 1, 2, 3, _, 5, 6, 7, _, 9]>
<MaskedTensor [_, 1, 2, 3, 4, _, 6, 7, 8, 9]>
<MaskedTensor [_, 1, 2, 3, 4, 5, _, 7, 8, 9]>

Batch e unbatching di set di dati con tipi di estensione

I set di dati con tipi di estensione possono essere batch e non batch utilizzando Dataset.batch e Dataset.unbatch .

batched_ds = ds.batch(2)
for value in batched_ds:
  print(value)

<MaskedTensor [[_, 1, _, 3, _, 5, _, 7, _, 9], [_, 1, 2, _, 4, 5, _, 7, 8, _]]>
<MaskedTensor [[_, 1, 2, 3, _, 5, 6, 7, _, 9], [_, 1, 2, 3, 4, _, 6, 7, 8, 9]]>
<MaskedTensor [[_, 1, 2, 3, 4, 5, _, 7, 8, 9]]>

unbatched_ds = batched_ds.unbatch()
for value in unbatched_ds:
  print(value)

<MaskedTensor [_, 1, _, 3, _, 5, _, 7, _, 9]>
<MaskedTensor [_, 1, 2, _, 4, 5, _, 7, 8, _]>
<MaskedTensor [_, 1, 2, 3, _, 5, 6, 7, _, 9]>
<MaskedTensor [_, 1, 2, 3, 4, _, 6, 7, 8, 9]>
<MaskedTensor [_, 1, 2, 3, 4, 5, _, 7, 8, 9]>