Ajuda a proteger a Grande Barreira de Corais com TensorFlow em Kaggle Junte Desafio

Tipos de extensão

Ver no TensorFlow.org Executar no Google Colab Ver fonte no GitHubBaixar caderno

Configurar

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

Tipos de extensão

Os tipos definidos pelo usuário podem tornar os projetos mais legíveis, modulares e de fácil manutenção. No entanto, a maioria das APIs do TensorFlow tem suporte muito limitado para tipos Python definidos pelo usuário. Isto inclui tanto APIs de alto nível (como Keras , tf.function , tf.SavedModel ) e de nível inferior APIs (tais como tf.while_loop e tf.concat ). Tipos de extensão TensorFlow pode ser usado para criar tipos definidos pelo usuário orientada a objetos que funcionam perfeitamente com APIs do TensorFlow. Para criar um tipo de extensão, simplesmente definir uma classe Python com tf.experimental.ExtensionType como a sua base, e usar anotações de tipo para especificar o tipo de cada 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

O tf.experimental.ExtensionType classe base funciona de forma semelhante ao typing.NamedTuple e @dataclasses.dataclass da biblioteca Python padrão. Em particular, ele adiciona automaticamente um construtor e métodos especiais (como __repr__ e __eq__ ) com base nas anotações de tipo de campo.

Normalmente, os tipos de extensão tendem a cair em uma de duas categorias:

  • As estruturas de dados, que agrupar um conjunto de valores relacionados, e pode proporcionar operações úteis com base nesses valores. As estruturas de dados pode ser bastante geral (tais como o TensorGraph exemplo acima); ou podem ser altamente personalizados para um modelo específico.

  • Tensor-like tipos, que se especializam ou estender o conceito de "Tensor." Tipos nesta categoria têm uma rank , uma shape , e geralmente um dtype ; e faz sentido para usá-los com operações tensor (tais como tf.stack , tf.add , ou tf.matmul ). MaskedTensor e CSRSparseMatrix são exemplos de tipos de tensor-like.

APIs com suporte

Os tipos de extensão são compatíveis com as seguintes APIs do TensorFlow:

  • Keras: tipos de extensão pode ser usado como entradas e saídas para Keras Models e Layers .
  • tf.data.Dataset: tipos de extensão podem ser incluídos em Datasets , e devolvido por conjunto de dados Iterators .
  • Concentrador Tensorflow: tipos de extensão pode ser utilizada como entradas e saídas para tf.hub módulos.
  • SavedModel: tipos de extensão pode ser usado como entradas e saídas para SavedModel funções.
  • tf.function: tipos de extensão podem ser usados como argumentos e valores de retorno para as funções envolvidas com o @tf.function decorador.
  • enquanto lacetes: tipos de extensão pode ser utilizado como variáveis de laço em tf.while_loop , e pode ser usado como argumentos e valores de retorno para o corpo de a, enquanto em malha.
  • condicionais: tipos de extensão podem ser condicionalmente selecionados usando tf.cond e tf.case .
  • py_function: tipos de extensão podem ser usados como argumentos e valores de retorno para o func argumento para tf.py_function .
  • Tensor ops: tipos de extensão podem ser estendidos para suportar a maioria das operações TensorFlow que aceita entradas tensor (por exemplo, tf.matmul , tf.gather , e tf.reduce_sum ). Consulte a seção "Despacho" abaixo para mais informações.
  • A estratégia de distribuição: tipos de extensão podem ser utilizados como valores por réplica.

Para obter mais detalhes, consulte a seção sobre "TensorFlow APIs que suportam ExtensionTypes" abaixo.

Requisitos

Tipos de campo

Todos os campos (também conhecidos como variáveis ​​de instância) devem ser declarados e uma anotação de tipo deve ser fornecida para cada campo. Os seguintes tipos de anotações são suportados:

Modelo Exemplo
Inteiros Python i: int
Python flutua f: float
Strings Python s: str
Booleanos Python b: bool
Python Nenhum n: None
Formas tensoras shape: tf.TensorShape
Tipos tensores dtype: tf.DType
Tensores t: tf.Tensor
Tipos de extensão mt: MyMaskedTensor
Tensores Ragged rt: tf.RaggedTensor
Tensores esparsos st: tf.SparseTensor
Fatias Indexadas s: tf.IndexedSlices
Tensores opcionais o: tf.experimental.Optional
Tipos de sindicatos int_or_float: typing.Union[int, float]
Tuplas params: typing.Tuple[int, float, tf.Tensor, int]
Var-length tuplas lengths: typing.Tuple[int, ...]
Mapeamentos tags: typing.Mapping[str, tf.Tensor]
Valores opcionais weight: typing.Optional[tf.Tensor]

Mutabilidade

Os tipos de extensão devem ser imutáveis. Isso garante que eles possam ser devidamente rastreados pelos mecanismos de rastreamento de gráfico do TensorFlow. Se você quiser alterar um valor de tipo de extensão, considere definir métodos que transformem os valores. Por exemplo, em vez de definir um set_mask método para transformar um MaskedTensor , você poderia definir um replace_mask método que retorna uma nova 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)

Funcionalidades adicionais por ExtensionType

O ExtensionType classe de base proporciona a seguinte funcionalidade:

  • Um construtor ( __init__ ).
  • Um método de representação de impressão ( __repr__ ).
  • Igualdade e operadores de desigualdade ( __eq__ ).
  • Um método de validação ( __validate__ ).
  • Imutabilidade forçada.
  • A aninhados TypeSpec .
  • Suporte para envio de API do Tensor.

Consulte a seção "Personalizando ExtensionTypes" abaixo para obter mais informações sobre como personalizar esta funcionalidade.

Construtor

O construtor adicionado por ExtensionType leva cada campo como um argumento nomeado (na ordem em que foram listados na definição de classe). Este construtor verificará o tipo de cada parâmetro e os converterá quando necessário. Em particular, Tensor campos são convertidos utilizando tf.convert_to_tensor ; Tuple campos são convertidos para tuple s; e Mapping campos são convertidos para dicts imutáveis.

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)

O construtor levanta uma TypeError se um valor campo não pode ser convertido para o seu tipo declarado:

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

O valor padrão de um campo pode ser especificado definindo seu valor no nível da 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>)

Representação para impressão

ExtensionType adiciona um método padrão de impressão representação ( __repr__ ) que inclui o nome da classe eo valor para cada 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])>)

Operadores de igualdade

ExtensionType acrescenta operadores de igualdade padrão ( __eq__ e __ne__ ) que consideram dois valores iguais se eles têm o mesmo tipo e todos os seus campos são iguais. Os campos do tensor são considerados iguais se tiverem a mesma forma e forem iguais para todos os elementos.

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

Método de validação

ExtensionType adiciona um __validate__ método, que pode ser substituído para executar verificações de validação em campos. Ele é executado depois que o construtor é chamado e depois que os campos são verificados e convertidos em seus tipos declarados, para que possa assumir que todos os campos têm seus tipos declarados.

ele seguinte exemplo atualiza MaskedTensor para validar a shape s e dtype s de seus campos:

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

Imutabilidade forçada

ExtensionType substitui o __setattr__ e __delattr__ métodos para prevenir a mutação, assegurando que os valores de extensão tipo são imutáveis.

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.

TypeSpec aninhado

Cada ExtensionType classe tem um correspondente TypeSpec classe, que é criado automaticamente e armazenado como <extension_type_name>.Spec .

Esta classe capta todas as informações a partir de um valor, exceto para os valores de quaisquer tensores aninhadas. Em particular, o TypeSpec para um valor é criado através da substituição de qualquer aninhada Tensor, ExtensionType, ou CompositeTensor com seu 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)})

TypeSpec valores podem ser construídos de forma explícita, ou eles podem ser construídos a partir de um ExtensionType valor usando tf.type_spec_from_value :

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

TypeSpec s são utilizadas por TensorFlow para dividir valores em um componente estático e um componente dinâmico:

  • O componente estático (que é fixado em tempo-gráfico construção) é codificado com um tf.TypeSpec .
  • A componente dinâmica (o qual pode variar a cada vez que o gráfico é executado) é codificado como uma lista de tf.Tensor s.

Por exemplo, tf.function refaz a sua função envolvido sempre que um argumento tem um inédito 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>}))

Para mais informações, consulte o tf.function Guia .

Personalizando ExtensionTypes

Além de simplesmente declarar campos e seus tipos, os tipos de extensão podem:

  • Substituir a representação de impressão padrão ( __repr__ ).
  • Defina métodos.
  • Defina métodos de classe e métodos estáticos.
  • Defina propriedades.
  • Substituir o construtor padrão ( __init__ ).
  • Substituir o operador de igualdade padrão ( __eq__ ).
  • Definir operadores (tais como __add__ e __lt__ ).
  • Declare valores padrão para os campos.
  • Defina subclasses.

Substituindo a representação padrão para impressão

Você pode substituir este operador de conversão de string padrão para tipos de extensão. O exemplo a seguir atualiza a MaskedTensor classe para gerar uma representação de cadeia mais legível quando os valores são impressos no modo ansioso.

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

Métodos de definição

Os tipos de extensão podem definir métodos, assim como qualquer classe Python normal. Por exemplo, o MaskedTensor tipo poderia definir um with_default método que retorna uma cópia de self com valores mascarados substituído por um determinado default de valor. Métodos podem, opcionalmente, ser anotado com o @tf.function decorador.

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

Definição de métodos de classe e métodos estáticos

Tipos de extensão podem definir métodos que utilizam os @classmethod e @staticmethod decoradores. Por exemplo, o MaskedTensor tipo poderia definir um método de fabricação que mascara qualquer elemento com um determinado valor:

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

Definindo propriedades

Tipos de extensão podem definir propriedades usando o @property decorador, como qualquer classe normais Python. Por exemplo, o MaskedTensor tipo poderia definir um dtype propriedade que é um atalho para o dtipo dos valores:

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

Substituindo o construtor padrão

Você pode substituir o construtor padrão para tipos de extensão. Os construtores personalizados devem definir um valor para cada campo declarado; e depois que o construtor personalizado retornar, todos os campos serão verificados e os valores serão convertidos conforme descrito acima.

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

Como alternativa, você pode considerar deixar o construtor padrão como está, mas adicionando um ou mais métodos de fábrica. Por exemplo:

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

Substituindo o operador de igualdade padrão ( __eq__ )

Você pode substituir o padrão __eq__ operador para tipos de extensão. O exemplo a seguir as atualizações MaskedTensor ignorar elementos mascarados quando se compara a igualdade.

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)

Usando referências futuras

Se o tipo de um campo ainda não foi definido, você pode usar uma string contendo o nome do tipo. No exemplo a seguir, a seqüência de "Node" é usado para anotar o children campo porque o Node tipo não foi (totalmente) definido ainda.

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

Definindo subclasses

Os tipos de extensão podem ser divididos em subclasses usando a sintaxe Python padrão. As subclasses de tipo de extensão podem adicionar novos campos, métodos e propriedades; e pode substituir o construtor, a representação para impressão e o operador de igualdade. O exemplo que se segue define uma base TensorGraph classe que utiliza três Tensor campos para codificar um conjunto de arestas entre os nós. Em seguida, ele define uma subclasse que adiciona um Tensor campo para gravar um "valor de recurso" para cada nó. A subclasse também define um método para propagar os valores do recurso ao longo das bordas.

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)

Definindo campos privados

Os campos de um tipo de extensão podem ser marcados como privados, prefixando-os com um sublinhado (seguindo as convenções padrão do Python). Isso não afeta a maneira como o TensorFlow trata os campos de forma alguma; mas simplesmente serve como um sinal para qualquer usuário do tipo de extensão de que esses campos são privados.

Personalização do ExtensionType TypeSpec

Cada ExtensionType classe tem um correspondente TypeSpec classe, que é criado automaticamente e armazenado como <extension_type_name>.Spec . Para obter mais informações, consulte a seção "TypeSpec aninhada" acima.

Para personalizar o TypeSpec , simplesmente definir sua própria classe aninhada chamada Spec , e ExtensionType vai usar isso como base para a construção automaticamente TypeSpec . Você pode personalizar o Spec classe por:

  • Substituindo a representação padrão para impressão.
  • Substituindo o construtor padrão.
  • Definição de métodos, métodos de classe, métodos estáticos e propriedades.

O exemplo a seguir personaliza a MaskedTensor.Spec classe para torná-lo mais fácil de usar:

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)

Envio da Tensor API

Tipos de extensão pode ser "tensor-like", no sentido de que eles são especializados ou estender a interface definida pelo tf.Tensor tipo. Exemplos de tipos de extensão de tensores semelhantes incluem RaggedTensor , SparseTensor , e MaskedTensor . Decoradores de despacho pode ser usado para substituir o comportamento padrão de operações TensorFlow quando aplicado a tensor-like tipos de extensão. O TensorFlow atualmente define três decoradores de despacho:

Despacho para uma única API

O tf.experimental.dispatch_for_api decorador substitui o comportamento padrão de uma operação TensorFlow especificado quando ele é chamado com a assinatura especificada. Por exemplo, você pode usar este decorador para especificar como tf.stack deve processar MaskedTensor valores:

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

Isso substitui a implementação padrão para tf.stack sempre que é chamado com uma lista de MaskedTensor valores (desde a values argumento é anotado com 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]]>

Para permitir tf.stack a listas alça de mistos MaskedTensor e Tensor valores, você pode refinar a anotação de tipo para o values de parâmetros e atualizar o corpo da função de forma adequada:

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

Para obter uma lista de APIs que pode ser substituído, consulte a documentação da API para tf.experimental.dispatch_for_api .

Despacho para todas as APIs elementwise unárias

O tf.experimental.dispatch_for_unary_elementwise_apis decorador substitui o comportamento padrão de todos os ops elemento a elemento unários (como tf.math.cos ) sempre que o valor para o primeiro argumento (geralmente chamado x ) corresponde ao tipo de anotação x_type . A função decorada deve ter dois argumentos:

  • api_func : Uma função que recebe um único parâmetro e executa a operação de elemento a elemento (por exemplo, tf.abs ).
  • x : O primeiro argumento para a operação de elemento a elemento.

O exemplo a seguir atualiza todas as operações elemento a elemento unários para lidar com a MaskedTensor Tipo:

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

Esta função irá agora ser usado sempre que uma operação de elemento a elemento unário é chamado em um 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]>

Envio para binários todas as APIs elementwise

Da mesma forma, tf.experimental.dispatch_for_binary_elementwise_apis pode ser usado para atualizar todas as operações elemento a elemento binários para lidar com a MaskedTensor Tipo:

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

Para obter uma lista das APIs elemento a elemento que são substituídas, consulte a documentação da API para tf.experimental.dispatch_for_unary_elementwise_apis e tf.experimental.dispatch_for_binary_elementwise_apis .

Batchable ExtensionTypes

Um ExtensionType é batchable se uma única instância pode ser usado para representar um grupo de valores. Tipicamente, isto é conseguido através da adição de dimensões de lote para todos aninhada Tensor s. As seguintes APIs do TensorFlow exigem que todas as entradas de tipo de extensão sejam batchable:

Por padrão, BatchableExtensionType cria valores em lote por dosadora de qualquer aninhada Tensor s, CompositeTensor s, e ExtensionType s. Se isso não é adequado para a sua classe, então você vai precisar usar tf.experimental.ExtensionTypeBatchEncoder para substituir esse comportamento padrão. Por exemplo, não seria adequado criar um lote de tf.SparseTensor valores simplesmente empilhamento tensores esparsas individuais values , indices e dense_shape campos - na maioria dos casos, você não pode empilhar estes tensores, uma vez que eles têm formas incompatíveis ; e mesmo se você pudesse, o resultado não seria um válido SparseTensor .

Exemplo de BatchableExtensionType: rede

Como exemplo, considere um simples Network de classe usado para balanceamento de carga, que rastreia a quantidade de trabalho que resta a fazer em cada nó, e quanta largura de banda está disponível para mover o trabalho entre os nós:

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

Para efectuar este tipo batchable, mudar o tipo de base para BatchableExtensionType , e ajustar a forma de cada campo de modo a incluir as dimensões de lote opcionais. O exemplo a seguir adiciona também uma shape campo para a faixa de keept a forma de lote. Esta shape campo não é obrigatório por tf.data.Dataset ou tf.map_fn , mas é exigido por 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.]]]>

Você pode então usar tf.data.Dataset para percorrer um lote de redes:

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 você também pode usar map_fn para aplicar uma função para cada elemento lote:

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

APIs TensorFlow compatíveis com ExtensionTypes

@ tf.function

tf.function é um decorador que precomputes TensorFlow gráficos para funções Python, que pode melhorar substancialmente o desempenho do seu código TensorFlow. Valores de tipo de extensão pode ser usado de forma transparente com @tf.function funções -decorated.

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 você deseja especificar explicitamente o input_signature para tf.function , então você pode fazer isso usando o tipo de extensão TypeSpec .

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

Funções concretas

Funções concretas encapsular gráficos traçados individuais que são construídas por tf.function . Os tipos de extensão podem ser usados ​​de forma transparente com funções concretas.

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

Operações de fluxo de controle

Os tipos de extensão são compatíveis com as operações de fluxo de controle do 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, _]>

Fluxo de controle de autógrafo

Os tipos de extensão também são suportados por instruções de fluxo de controle em tf.function (usando autógrafo). No exemplo a seguir, o if declaração e for as declarações são automaticamente convertidos para tf.cond e tf.while_loop operações, que tipos de extensão de suporte.

@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 é de TensorFlow API de alto nível para construir e treinar modelos de aprendizagem profunda. Os tipos de extensão podem ser passados ​​como entradas para um modelo Keras, passados ​​entre as camadas Keras e retornados pelos modelos Keras. Keras atualmente coloca dois requisitos nos tipos de extensão:

  • Eles devem ser batchable (consulte "Batchable ExtensionTypes" acima).
  • A deve ter um campo ou propriedade denominada shape . shape[0] é assumido como sendo a dimensão do lote.

As duas subseções a seguir fornecem exemplos que mostram como os tipos de extensão podem ser usados ​​com o Keras.

Keras exemplo: Network

Para o primeiro exemplo, considere a Network classe definida na seção "Batchable ExtensionTypes" acima, que pode ser usado para balanceamento de carga de trabalho entre os nós. Sua definição é repetida aqui:

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

Você pode definir uma nova camada Keras que processa Network s.

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)

Você pode então usar essas camadas para criar um modelo simples. Para alimentar um ExtensionType em um modelo, você pode usar um tf.keras.layer.Input camada com type_spec conjunto para o tipo de extensão TypeSpec . Se o modelo Keras vai ser utilizado para processar lotes, em seguida, o type_spec deve incluir a dimensão do lote.

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

Finalmente, você pode aplicar o modelo a uma única rede e a um lote de redes.

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

Exemplo de Keras: MaskedTensor

Neste exemplo, MaskedTensor é estendido para suportar Keras . shape é definido como uma propriedade que é calculado a partir da values de campo. Keras requer thatyou adicionar essa propriedade para ambos o tipo de extensão e sua TypeSpec . MaskedTensor também define um __name__ variável, que será necessário para SavedModel serialização (abaixo).

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

Em seguida, os decoradores de despacho são usados ​​para substituir o comportamento padrão de várias APIs do TensorFlow. Uma vez que estas API são utilizados por camadas Keras convencionais (tais como a Dense camada), sobrepondo-se estes nos permitirá usar essas camadas com MaskedTensor . Para os fins deste exemplo, matmul para tensores mascarados é definido para tratar os valores mascarados como zeros (ou seja, a não incluí-los no produto).

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

Você pode então construir um modelo Keras que aceita MaskedTensor entradas, usando camadas Keras padrão:

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)

SavedModel

Um SavedModel é um programa TensorFlow em série, incluindo ambos os pesos e cálculo. Ele pode ser construído a partir de um modelo Keras ou de um modelo personalizado. Em ambos os casos, os tipos de extensão podem ser usados ​​de forma transparente com as funções e métodos definidos por um SavedModel.

SavedModel pode salvar modelos, camadas e funções que tipos de extensão processo, desde que os tipos de extensão têm um __name__ campo. Este nome é usado para registrar o tipo de extensão, para que possa ser localizado quando o modelo for carregado.

Exemplo: salvar um modelo Keras

Modelos Keras que usam tipos de extensão podem ser salvas usando 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)>

Exemplo: salvar um modelo personalizado

SavedModel também pode ser usado para salvar costume tf.Module subclasses com funções que tipos de extensão processo.

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

Carregando um SavedModel quando ExtensionType não está disponível

Se você carregar um SavedModel que usa um ExtensionType , mas que ExtensionType não está disponível (ou seja, não foi importado), então você verá um aviso e TensorFlow vai cair para trás a usar um objeto "anonymous tipo de extensão". Este objeto terá os mesmos campos do tipo original, mas não terá nenhuma personalização adicional que você tenha adicionado para o tipo, como métodos ou propriedades personalizadas.

Usando ExtensionTypes com serviço TensorFlow

Atualmente, TensorFlow servindo (e outros consumidores do SavedModel "assinaturas" dicionário) exigem que todas as entradas e saídas de ser tensores matérias. Se quiser usar o TensorFlow servindo com um modelo que usa tipos de extensão, você pode adicionar métodos de empacotador que compõem ou decompõem valores de tipo de extensão de tensores. Por exemplo:

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>

Conjuntos de dados

tf.data é uma API que permite que você construa dutos de entrada complexos a partir de simples, peças reutilizáveis. A sua estrutura de dados de núcleo é tf.data.Dataset , que representa uma sequência de elementos, em que cada elemento é constituído de um ou mais componentes.

Construindo conjuntos de dados com tipos de extensão

Conjuntos de dados podem ser construídos a partir de valores de tipo de extensão usando Dataset.from_tensors , Dataset.from_tensor_slices , ou 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]>

Colocação em lote e desencadeamento de conjuntos de dados com tipos de extensão

Conjuntos de dados com tipos de extensão pode ser batchand e unbatched usando 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]>