Rodzaje rozszerzeń

Zobacz na TensorFlow.org Uruchom w Google Colab Wyświetl źródło na GitHubPobierz notatnik

Ustawiać

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

Rodzaje rozszerzeń

Typy zdefiniowane przez użytkownika mogą sprawić, że projekty będą bardziej czytelne, modułowe i łatwe w utrzymaniu. Jednak większość interfejsów API TensorFlow ma bardzo ograniczoną obsługę typów Pythona zdefiniowanych przez użytkownika. Obejmuje to zarówno interfejsy API wysokiego poziomu (takie jak Keras , tf.function , tf.SavedModel ) jak i interfejsy API niższego poziomu (takie jak tf.while_loop i tf.concat ). Typy rozszerzeń TensorFlow mogą być używane do tworzenia zdefiniowanych przez użytkownika typów obiektowych, które bezproblemowo współpracują z interfejsami API TensorFlow. Aby utworzyć typ rozszerzenia, po prostu zdefiniuj klasę Pythona z tf.experimental.ExtensionType jako jej podstawą i użyj adnotacji typu , aby określić typ dla każdego pola.

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

Klasa bazowa tf.experimental.ExtensionType działa podobnie do typing.NamedTuple i @dataclasses.dataclass ze standardowej biblioteki Pythona. W szczególności automatycznie dodaje konstruktor i metody specjalne (takie jak __repr__ i __eq__ ) na podstawie adnotacji typu pola.

Zazwyczaj typy rozszerzeń należą do jednej z dwóch kategorii:

  • Struktury danych , które grupują zbiór powiązanych wartości i mogą udostępniać przydatne operacje oparte na tych wartościach. Struktury danych mogą być dość ogólne (takie jak powyższy przykład TensorGraph ); lub mogą być wysoce dostosowane do konkretnego modelu.

  • Typy tensorowe , które specjalizują się lub rozszerzają pojęcie „Tensor”. Typy w tej kategorii mają rank , shape i zwykle dtype ; i sensowne jest ich używanie z operacjami Tensor (takimi jak tf.stack , tf.add lub tf.matmul ). MaskedTensor i CSRSparseMatrix to przykłady typów podobnych do tensora.

Obsługiwane interfejsy API

Typy rozszerzeń są obsługiwane przez następujące interfejsy API TensorFlow:

  • Keras : Typy rozszerzeń mogą być używane jako dane wejściowe i wyjściowe dla Models i Layers Keras.
  • tf.data.Dataset : typy rozszerzeń mogą być zawarte w zestawach danych i zwracane przez Iterators Datasets danych .
  • Tensorflow hub : typy rozszerzeń mogą być używane jako wejścia i wyjścia dla modułów tf.hub .
  • SavedModel : Typy rozszerzeń mogą być używane jako wejścia i wyjścia dla funkcji SavedModel .
  • tf.function : typy rozszerzeń mogą być używane jako argumenty i wartości zwracane dla funkcji opakowanych w dekorator @tf.function .
  • pętle while : typy rozszerzeń mogą być używane jako zmienne pętli w tf.while_loop i mogą być używane jako argumenty i wartości zwracane dla treści pętli while.
  • warunkowe : Typy rozszerzeń można warunkowo wybierać za pomocą tf.cond i tf.case .
  • py_function : typy rozszerzeń mogą być używane jako argumenty i zwracane wartości dla argumentu func do tf.py_function .
  • Tensor ops : typy rozszerzeń można rozszerzyć, aby obsługiwały większość operacji TensorFlow, które akceptują dane wejściowe Tensor (np. tf.matmul , tf.gather i tf.reduce_sum ). Zobacz sekcję „ Wysyłka ” poniżej, aby uzyskać więcej informacji.
  • strategia dystrybucji : typy rozszerzeń mogą być używane jako wartości na replikę.

Aby uzyskać więcej informacji, zobacz sekcję „Interfejsy API TensorFlow obsługujące typy rozszerzeń” poniżej.

Wymagania

Rodzaje pól

Wszystkie pola (inaczej zmienne instancji) muszą być zadeklarowane, a dla każdego pola należy podać adnotację typu. Obsługiwane są następujące adnotacje typu:

Rodzaj Przykład
Liczby całkowite w Pythonie i: int
Python unosi się na wodzie f: float
Łańcuchy Pythona s: str
Wartości logiczne Pythona b: bool
Python Brak n: None
Kształty tensorów shape: tf.TensorShape
Typy tensorów dtype: tf.DType
Tensory t: tf.Tensor
Rodzaje rozszerzeń mt: MyMaskedTensor
Poszarpane Tensory rt: tf.RaggedTensor
Rzadkie Tensory st: tf.SparseTensor
Zindeksowane plastry s: tf.IndexedSlices
Opcjonalne Tensory o: tf.experimental.Optional
Wpisz związki int_or_float: typing.Union[int, float]
Krotki params: typing.Tuple[int, float, tf.Tensor, int]
Krotki o długości var lengths: typing.Tuple[int, ...]
Mapowania tags: typing.Mapping[str, tf.Tensor]
Wartości opcjonalne weight: typing.Optional[tf.Tensor]

Zmienność

Typy rozszerzeń muszą być niezmienne. Gwarantuje to, że mogą być odpowiednio śledzone przez mechanizmy śledzenia wykresów TensorFlow. Jeśli chcesz zmutować wartość typu rozszerzenia, rozważ zamiast tego zdefiniowanie metod, które przekształcają wartości. Na przykład, zamiast definiować metodę set_mask do mutacji MaskedTensor , można zdefiniować metodę replace_mask , która zwraca nowy 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)

Funkcjonalność dodana przez ExtensionType

Klasa bazowa ExtensionType zapewnia następującą funkcjonalność:

  • Konstruktor ( __init__ ).
  • Drukowalna metoda reprezentacji ( __repr__ ).
  • Operatory równości i nierówności ( __eq__ ).
  • Metoda weryfikacji ( __validate__ ).
  • Wymuszona niezmienność.
  • Zagnieżdżony TypeSpec .
  • Obsługa wysyłki Tensor API.

Zobacz sekcję „Dostosowywanie typów rozszerzeń” poniżej, aby uzyskać więcej informacji na temat dostosowywania tej funkcji.

Konstruktor

Konstruktor dodany przez ExtensionType przyjmuje każde pole jako nazwany argument (w kolejności, w jakiej zostały wymienione w definicji klasy). Ten konstruktor sprawdzi każdy parametr i przekonwertuje go w razie potrzeby. W szczególności pola Tensor są konwertowane za pomocą tf.convert_to_tensor ; Pola Tuple są konwertowane na tuple ; Pola Mapping są konwertowane na niezmienne wersety.

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)

Konstruktor TypeError , jeśli wartości pola nie można przekonwertować na zadeklarowany typ:

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

Domyślną wartość pola można określić, ustawiając jego wartość na poziomie klasy:

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>)

Reprezentacja do druku

ExtensionType dodaje domyślną, drukowalną metodę reprezentacji ( __repr__ ), która zawiera nazwę klasy i wartość każdego pola:

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])>)

Operatorzy równości

ExtensionType dodaje domyślne operatory równości ( __eq__ i __ne__ ), które uznają dwie wartości za równe, jeśli mają ten sam typ i wszystkie ich pola są równe. Pola tensorowe są uważane za równe, jeśli mają ten sam kształt i są równe elementowo dla wszystkich elementów.

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

Metoda walidacji

ExtensionType dodaje metodę __validate__ , którą można nadpisać w celu sprawdzenia poprawności pól. Jest uruchamiany po wywołaniu konstruktora i po sprawdzeniu typu pól i przekonwertowaniu ich na zadeklarowane typy, dzięki czemu można założyć, że wszystkie pola mają zadeklarowane typy.

Poniższy przykład aktualizuje MaskedTensor , aby sprawdzić poprawność shape s i dtype s jego pól:

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

Wymuszona niezmienność

ExtensionType zastępuje metody __setattr__ i __delattr__ , aby zapobiec mutacji, zapewniając, że wartości typu rozszerzenia są niezmienne.

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.

Specyfikacja typu zagnieżdżonego

Każda klasa ExtensionType ma odpowiadającą jej klasę TypeSpec , która jest tworzona automatycznie i przechowywana jako <extension_type_name>.Spec .

Ta klasa przechwytuje wszystkie informacje z wartości z wyjątkiem wartości dowolnych zagnieżdżonych tensorów. W szczególności TypeSpec dla wartości jest tworzony przez zastąpienie dowolnego zagnieżdżonego Tensor, ExtensionType lub CompositeTensor jego 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)})

Wartości TypeSpec mogą być konstruowane jawnie lub mogą być budowane z wartości ExtensionType przy użyciu tf.type_spec_from_value :

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

TypeSpec są używane przez TensorFlow do dzielenia wartości na składnik statyczny i składnik dynamiczny :

  • Składnik statyczny (który jest ustalany w czasie tworzenia wykresu) jest kodowany za pomocą tf.TypeSpec .
  • Komponent dynamiczny (który może się zmieniać przy każdym uruchomieniu wykresu) jest zakodowany jako lista tf.Tensor s.

Na przykład tf.function odtwarza swoją opakowaną funkcję za każdym razem, gdy argument ma wcześniej niewidoczną TypeSpec :

@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>}))

Więcej informacji można znaleźć w przewodniku tf.function Guide .

Dostosowywanie typów rozszerzeń

Oprócz prostego deklarowania pól i ich typów, typy rozszerzeń mogą:

  • Zastąp domyślną reprezentację drukowalną ( __repr__ ).
  • Zdefiniuj metody.
  • Zdefiniuj metody klasowe i statyczne.
  • Zdefiniuj właściwości.
  • Zastąp domyślny konstruktor ( __init__ ).
  • Zastąp domyślny operator równości ( __eq__ ).
  • Zdefiniuj operatory (takie jak __add__ i __lt__ ).
  • Zadeklaruj wartości domyślne dla pól.
  • Zdefiniuj podklasy.

Zastępowanie domyślnej reprezentacji do druku

Możesz zastąpić ten domyślny operator konwersji ciągów dla typów rozszerzeń. Poniższy przykład aktualizuje klasę MaskedTensor , aby wygenerować bardziej czytelną reprezentację ciągu, gdy wartości są drukowane w trybie 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]]>

Definiowanie metod

Typy rozszerzeń mogą definiować metody, tak jak każda normalna klasa Pythona. Na przykład typ MaskedTensor może zdefiniować metodę with_default , która zwraca kopię self z zamaskowanymi wartościami zastąpionymi podaną wartością default . Metody mogą opcjonalnie być opatrzone adnotacjami za pomocą dekoratora @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)>

Definiowanie metod klasowych i statycznych

Typy rozszerzeń mogą definiować metody przy użyciu dekoratorów @classmethod i @staticmethod . Na przykład typ MaskedTensor mógłby zdefiniować metodę fabryki, która maskuje dowolny element o podanej wartości:

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]]>

Definiowanie właściwości

Typy rozszerzeń mogą definiować właściwości za pomocą dekoratora @property , tak jak każda normalna klasa Pythona. Na przykład typ MaskedTensor może definiować właściwość dtype , która jest skrótem dla dtype wartości:

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

Zastępowanie domyślnego konstruktora

Możesz zastąpić domyślny konstruktor dla typów rozszerzeń. Konstruktory niestandardowe muszą ustawić wartość dla każdego zadeklarowanego pola; a po powrocie konstruktora niestandardowego wszystkie pola zostaną sprawdzone pod względem typu, a wartości zostaną przekonwertowane zgodnie z powyższym opisem.

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>)

Alternatywnie możesz rozważyć pozostawienie domyślnego konstruktora bez zmian, ale dodanie jednej lub więcej metod fabrycznych. Np:

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>)

Zastępowanie domyślnego operatora równości ( __eq__ )

Możesz zastąpić domyślny operator __eq__ dla typów rozszerzeń. Poniższy przykład aktualizuje MaskedTensor , aby ignorować zamaskowane elementy podczas porównywania pod kątem równości.

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)

Korzystanie z odwołań do przodu

Jeśli typ pola nie został jeszcze zdefiniowany, możesz zamiast niego użyć ciągu zawierającego nazwę typu. W poniższym przykładzie ciąg "Node" jest używany do children pola child, ponieważ typ Node nie został jeszcze (w pełni) zdefiniowany.

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=())))

Definiowanie podklas

Typy rozszerzeń mogą być podzielone na podklasy przy użyciu standardowej składni Pythona. Podklasy typu rozszerzenia mogą dodawać nowe pola, metody i właściwości; i może przesłonić konstruktor, reprezentację drukowalną i operator równości. Poniższy przykład definiuje podstawową klasę TensorGraph , która używa trzech pól Tensor do kodowania zestawu krawędzi między węzłami. Następnie definiuje podklasę, która dodaje pole Tensor , aby zarejestrować „wartość cechy” dla każdego węzła. Podklasa definiuje również metodę propagowania wartości cech wzdłuż krawędzi.

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)

Definiowanie pól prywatnych

Pola typu rozszerzenia można oznaczyć jako prywatne, poprzedzając je podkreśleniem (zgodnie ze standardowymi konwencjami Pythona). Nie ma to wpływu na sposób, w jaki TensorFlow traktuje pola w jakikolwiek sposób; ale po prostu służy jako sygnał dla wszystkich użytkowników typu rozszerzenia, że ​​te pola są prywatne.

Dostosowywanie specyfikacji TypeSpec rozszerzenia ExtensionType

Każda klasa ExtensionType ma odpowiadającą jej klasę TypeSpec , która jest tworzona automatycznie i przechowywana jako <extension_type_name>.Spec . Aby uzyskać więcej informacji, zobacz sekcję „Zagnieżdżona specyfikacja typów” powyżej.

Aby dostosować TypeSpec , po prostu zdefiniuj własną zagnieżdżoną klasę o nazwie Spec , a ExtensionType użyje jej jako podstawy dla automatycznie skonstruowanej TypeSpec . Możesz dostosować klasę Spec poprzez:

  • Zastępowanie domyślnej reprezentacji do druku.
  • Zastępowanie domyślnego konstruktora.
  • Definiowanie metod, metod klasowych, metod statycznych i właściwości.

Poniższy przykład dostosowuje klasę MaskedTensor.Spec , aby ułatwić jej użycie:

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)

Wysyłka API Tensora

Typy rozszerzeń mogą być „podobne do tensora”, w tym sensie, że specjalizują się lub rozszerzają interfejs zdefiniowany przez typ tf.Tensor . Przykłady typów rozszerzeń podobnych do tensora obejmują RaggedTensor , SparseTensor i MaskedTensor . Dekoratory wysyłania mogą służyć do zastępowania domyślnego zachowania operacji TensorFlow w przypadku zastosowania do typów rozszerzeń podobnych do tensorów. TensorFlow obecnie definiuje trzy dekoratory wysyłek:

Wysyłka dla jednego API

Dekorator tf.experimental.dispatch_for_api zastępuje domyślne zachowanie określonej operacji TensorFlow, gdy jest wywoływana z określoną sygnaturą. Na przykład możesz użyć tego dekoratora, aby określić, jak tf.stack powinien przetwarzać wartości 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))

Zastępuje to domyślną implementację dla tf.stack za każdym razem, gdy jest wywoływana z listą wartości MaskedTensor (ponieważ argument values jest oznaczony 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]]>

Aby umożliwić tf.stack obsługę list mieszanych wartości MaskedTensor i Tensor , można udoskonalić adnotację typu dla parametru values i odpowiednio zaktualizować treść funkcji:

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, _]]>

Listę interfejsów API, które można zastąpić, można znaleźć w dokumentacji interfejsu API dla tf.experimental.dispatch_for_api .

Wysyłka dla wszystkich jednoargumentowych interfejsów API elementwise

Dekorator tf.experimental.dispatch_for_unary_elementwise_apis przesłania domyślne zachowanie wszystkich operacji jednoargumentowych elementwise (takich jak tf.math.cos ), gdy wartość pierwszego argumentu (zazwyczaj o nazwie x ) pasuje do adnotacji typu x_type . Zdobiona funkcja powinna mieć dwa argumenty:

  • api_func : Funkcja, która przyjmuje pojedynczy parametr i wykonuje operację elementwise (np. tf.abs ).
  • x : pierwszy argument operacji elementarnej.

Poniższy przykład aktualizuje wszystkie jednoargumentowe operacje elementwise w celu obsługi typu 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)

Ta funkcja będzie teraz używana za każdym razem, gdy jednoargumentowa operacja elementowa jest wywoływana na 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]>

Wyślij binarne wszystkie API elementwise

Podobnie, tf.experimental.dispatch_for_binary_elementwise_apis może służyć do aktualizacji wszystkich binarnych operacji elementwise do obsługi typu 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], [_, _, _]]>

Listę interfejsów API elementwise, które są zastępowane, można znaleźć w dokumentacji interfejsu API dla tf.experimental.dispatch_for_unary_elementwise_apis i tf.experimental.dispatch_for_binary_elementwise_apis .

Typy rozszerzeń wsadowych

ExtensionType można przetwarzać wsadowo , jeśli pojedyncze wystąpienie może służyć do reprezentowania partii wartości. Zazwyczaj osiąga się to poprzez dodanie wymiarów wsadowych do wszystkich zagnieżdżonych Tensor . Następujące interfejsy API TensorFlow wymagają, aby wszelkie dane wejściowe typu rozszerzenia były grupowe:

Domyślnie BatchableExtensionType tworzy wartości wsadowe, grupując dowolne zagnieżdżone Tensor s, CompositeTensor s i ExtensionType s. Jeśli nie jest to odpowiednie dla Twojej klasy, będziesz musiał użyć tf.experimental.ExtensionTypeBatchEncoder , aby zastąpić to domyślne zachowanie. Na przykład nie byłoby właściwe tworzenie partii wartości tf.SparseTensor przez proste zestawienie poszczególnych values tensorów sparse , indices i pól dense_shape — w większości przypadków nie można układać tych tensorów w stos, ponieważ mają one niezgodne kształty ; a nawet gdybyś mógł, wynik nie byłby prawidłowym SparseTensor .

Przykład BatchableExtensionType: Sieć

Jako przykład rozważmy prostą klasę Network używaną do równoważenia obciążenia, która śledzi, ile pracy pozostało do wykonania w każdym węźle i jaka przepustowość jest dostępna do przenoszenia pracy między węzłami:

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]])

Aby uczynić ten typ partiami, zmień typ podstawowy na BatchableExtensionType i dostosuj kształt każdego pola, aby uwzględnić opcjonalne wymiary partii. Poniższy przykład dodaje również pole shape , aby śledzić kształt partii. To pole shape nie jest wymagane przez tf.data.Dataset ani tf.map_fn , ale jest wymagane przez 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.]]]>

Następnie możesz użyć tf.data.Dataset , aby przejść przez partię sieci:

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.]]>

Możesz także użyć map_fn , aby zastosować funkcję do każdego elementu wsadowego:

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.]]]>

Interfejsy API TensorFlow obsługujące typy rozszerzeń

@tf.funkcja

tf.function to dekorator, który wstępnie oblicza wykresy TensorFlow dla funkcji Pythona, co może znacznie poprawić wydajność Twojego kodu TensorFlow. Wartości typu rozszerzenia mogą być używane w sposób przezroczysty z @tf.function .

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)>

Jeśli chcesz jawnie określić input_signature dla tf.function , możesz to zrobić za pomocą TypeSpec typu rozszerzenia.

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)>)

Konkretne funkcje

Konkretne funkcje hermetyzują poszczególne grafy śledzone, które są budowane przez tf.function . Typy rozszerzeń mogą być używane w sposób przezroczysty z konkretnymi funkcjami.

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

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

Kontroluj operacje przepływu

Typy rozszerzeń są obsługiwane przez operacje przepływu sterowania 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, _]>

Kontrola przebiegu autografu

Typy rozszerzeń są również obsługiwane przez instrukcje przepływu sterowania w tf.function (przy użyciu autografu). W poniższym przykładzie instrukcje if i instrukcje for są automatycznie konwertowane na tf.cond i tf.while_loop , które obsługują typy rozszerzeń.

@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]>

Keras

tf.keras to wysokopoziomowy interfejs API firmy TensorFlow do tworzenia i trenowania modeli uczenia głębokiego. Typy rozszerzeń mogą być przekazywane jako dane wejściowe do modelu Keras, przekazywane między warstwami Keras i zwracane przez modele Keras. Keras stawia obecnie dwa wymagania dotyczące typów rozszerzeń:

  • Muszą nadawać się do przetwarzania wsadowego (patrz „Batchable ExtensionTypes” powyżej).
  • Musi mieć pole lub właściwość o nazwie shape . Przyjmuje się, że shape[0] jest wymiarem partii.

Poniższe dwie podsekcje zawierają przykłady pokazujące, w jaki sposób typy rozszerzeń mogą być używane z Keras.

Przykład Kerasa: Network

W pierwszym przykładzie rozważmy klasę Network zdefiniowaną w sekcji „Batchable ExtensionTypes” powyżej, która może służyć do równoważenia obciążenia między węzłami. Jego definicja jest tutaj powtórzona:

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]]])

Możesz zdefiniować nową warstwę Keras, która przetwarza 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)

Następnie możesz użyć tych warstw do stworzenia prostego modelu. Aby wprowadzić ExtensionType do modelu, można użyć warstwy tf.keras.layer.Input z type_spec ustawionym na TypeSpec typu rozszerzenia. Jeśli model Keras będzie używany do przetwarzania partii, wówczas parametr type_spec musi zawierać wymiar partii.

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(),
    ])

Na koniec możesz zastosować model do pojedynczej sieci i do partii sieci.

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.]]]>

Przykład Keras: MaskedTensor

W tym przykładzie MaskedTensor jest rozszerzony o obsługę Keras . shape jest definiowany jako właściwość obliczana z pola values . Keras wymaga dodania tej właściwości zarówno do typu rozszerzenia, jak i do jego TypeSpec . MaskedTensor definiuje również zmienną __name__ , która będzie wymagana do serializacji SavedModel (poniżej).

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))

Następnie dekoratory wysyłania są używane do zastąpienia domyślnego zachowania kilku interfejsów API TensorFlow. Ponieważ te interfejsy API są używane przez standardowe warstwy Keras (takie jak warstwa Dense ), zastąpienie ich pozwoli nam używać tych warstw z MaskedTensor . Na potrzeby tego przykładu, matmul dla zamaskowanych tensorów jest zdefiniowany w celu traktowania zamaskowanych wartości jako zer (tj. aby nie uwzględniać ich w produkcie).

@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)

Następnie można skonstruować model Keras, który akceptuje dane wejściowe MaskedTensor , używając standardowych warstw Keras:

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)

Zapisany model

SavedModel to serializowany program TensorFlow, obejmujący zarówno wagi, jak i obliczenia. Może być zbudowany z modelu Keras lub z modelu niestandardowego. W obu przypadkach typy rozszerzeń mogą być używane w sposób przezroczysty z funkcjami i metodami zdefiniowanymi przez SavedModel.

SavedModel może zapisywać modele, warstwy i funkcje, które przetwarzają typy rozszerzeń, o ile typy rozszerzeń mają pole __name__ . Ta nazwa służy do rejestrowania typu rozszerzenia, dzięki czemu można ją zlokalizować, gdy model jest ładowany.

Przykład: zapisywanie modelu Keras

Modele Keras, które używają typów rozszerzeń, można zapisać za pomocą 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)>

Przykład: zapisywanie niestandardowego modelu

SavedModel może być również używany do zapisywania niestandardowych podklas tf.Module z funkcjami przetwarzającymi typy rozszerzeń.

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, _]>

Ładowanie SavedModel, gdy ExtensionType jest niedostępny

Jeśli załadujesz SavedModel , który używa ExtensionType , ale ten ExtensionType nie jest dostępny (tj. nie został zaimportowany), zobaczysz ostrzeżenie, a TensorFlow wróci do używania obiektu „anonimowego typu rozszerzenia”. Ten obiekt będzie miał te same pola, co oryginalny typ, ale nie będzie w nim żadnych dalszych dostosowań dodanych dla typu, takich jak niestandardowe metody lub właściwości.

Korzystanie z typów rozszerzeń z obsługą TensorFlow

Obecnie obsługa TensorFlow (i inni użytkownicy słownika „sygnatur” SavedModel) wymagają, aby wszystkie dane wejściowe i wyjściowe były surowymi tensorami. Jeśli chcesz korzystać z udostępniania TensorFlow z modelem, który używa typów rozszerzeń, możesz dodać metody opakowujące, które komponują lub rozkładają wartości typów rozszerzeń z tensorów. Np:

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>

Zbiory danych

tf.data to interfejs API, który umożliwia budowanie złożonych potoków wejściowych z prostych elementów wielokrotnego użytku. Jego podstawową strukturą danych jest tf.data.Dataset , która reprezentuje sekwencję elementów, w której każdy element składa się z jednego lub więcej składników.

Budowanie zbiorów danych z typami rozszerzeń

Zestawy danych można budować z wartości typu rozszerzenia za pomocą Dataset.from_tensors , Dataset.from_tensor_slices lub 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]>

Grupowanie i rozdzielanie zbiorów danych z typami rozszerzeń

Zestawy danych z typami rozszerzeń mogą być wsadowe i niewsadowe przy użyciu Dataset.batch adn 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]>