ヘルプKaggleにTensorFlowグレートバリアリーフを保護チャレンジに参加

拡張タイプ

TensorFlow.orgで表示GoogleColabで実行GitHubでソースを表示ノートブックをダウンロード

設定

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

拡張タイプ

ユーザー定義型は、プロジェクトをより読みやすく、モジュール式で、保守しやすくすることができます。ただし、ほとんどのTensorFlow APIは、ユーザー定義のPython型のサポートが非常に限られています。これには、高レベルAPI( Kerastf.functiontf.SavedModelなど)と低レベルAPI( tf.while_looptf.concatなど)の両方が含まれます。 TensorFlow拡張タイプを使用して、TensorFlowのAPIとシームレスに連携するユーザー定義のオブジェクト指向タイプを作成できます。拡張型を作成するには、 tf.experimental.ExtensionTypeをベースとしてPythonクラスを定義し、型注釈を使用して各フィールドの型を指定するだけです。

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

tf.experimental.ExtensionType基本クラスは、標準のPythonライブラリのtyping.NamedTupleおよび@dataclasses.dataclassと同様に機能します。特に、フィールドタイプの注釈に基づいて、コンストラクターと特別なメソッド( __repr____eq__など)を自動的に追加します。

通常、拡張タイプは次の2つのカテゴリのいずれかに分類される傾向があります。

  • 関連する値のコレクションをグループ化し、それらの値に基づいて有用な操作を提供できるデータ構造。データ構造はかなり一般的である可能性があります(上記のTensorGraphの例など)。または、特定のモデルに合わせて高度にカスタマイズすることもできます。

  • 「テンソル」の概念を専門化または拡張するテンソルのようなタイプ。このカテゴリのタイプには、 rankshape 、および通常はdtypeがあります。そして、それらをTensor操作( tf.stacktf.addtf.matmulなど)で使用することは理にかなっています。 MaskedTensorCSRSparseMatrixは、テンソルのようなタイプの例です。

サポートされているAPI

拡張タイプは、次のTensorFlowAPIでサポートされています。

  • Keras :拡張タイプはKeras ModelsLayersの入力と出力として使用できます。
  • tf.data.Dataset :拡張タイプはDatasetsに含めることができ、データセットIteratorsによって返されます。
  • Tensorflowハブ:拡張タイプは、 tf.hubモジュールの入力と出力として使用できます。
  • SavedModel :拡張タイプは、 SavedModel関数の入力および出力として使用できます。
  • tf.function :拡張タイプは、 @tf.functionデコレータでラップされた関数の引数および戻り値として使用できます。
  • whileループ:拡張タイプは、 tf.while_loopのループ変数として使用でき、whileループの本体の引数および戻り値として使用できます。
  • 条件付き:拡張タイプは、 tf.condおよびtf.caseを使用して条件付きで選択できます。
  • py_function :拡張タイプは引数として使用でき、 tf.py_functionへのfunc引数の値を返します。
  • Tensor ops :拡張タイプを拡張して、Tensor入力を受け入れるほとんどのTensorFlow opsをサポートできます(例: tf.matmultf.gathertf.reduce_sum )。詳細については、以下の「ディスパッチ」セクションを参照してください。
  • 配布戦略:拡張タイプは、レプリカごとの値として使用できます。

詳細については、以下の「ExtensionTypesをサポートするTensorFlowAPI」のセクションをご覧ください。

要件

フィールドタイプ

すべてのフィールド(別名インスタンス変数)を宣言する必要があり、フィールドごとに型注釈を指定する必要があります。次の型注釈がサポートされています。

タイプ
Python整数i: int
Pythonフロートf: float
Python文字列s: str
Pythonブール値b: bool
Pythonなしn: None
テンソル形状shape: tf.TensorShape
テンソルdtypes dtype: tf.DType
テンソルt: tf.Tensor
拡張タイプmt: MyMaskedTensor
不規則テンソルrt: tf.RaggedTensor
スパーステンソルst: tf.SparseTensor
インデックス付きスライスs: tf.IndexedSlices
オプションのテンソルo: tf.experimental.Optional
タイプユニオンint_or_float: typing.Union[int, float]
タプルparams: typing.Tuple[int, float, tf.Tensor, int]
Varの長さのタプルlengths: typing.Tuple[int, ...]
マッピングtags: typing.Mapping[str, tf.Tensor]
オプション値weight: typing.Optional[tf.Tensor]

可変性

拡張タイプは不変である必要があります。これにより、TensorFlowのグラフトレースメカニズムによって適切に追跡できるようになります。拡張タイプの値を変更したい場合は、代わりに値を変換するメソッドを定義することを検討してください。たとえば、MaskedTensorを変更するためにset_maskメソッドを定義するのではなく、新しいMaskedTensorを返すreplace_maskメソッドを定義でき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)

ExtensionTypeによって追加された機能

ExtensionType基本クラスは、次の機能を提供します。

  • コンストラクター( __init__ )。
  • 印刷可能な表現方法( __repr__ )。
  • 等式および不等式演算子( __eq__ )。
  • 検証方法( __validate__ )。
  • 強制された不変性。
  • ネストされたTypeSpec
  • TensorAPIディスパッチのサポート。

この機能のカスタマイズの詳細については、以下の「ExtensionTypesのカスタマイズ」セクションを参照してください。

コンストラクタ

ExtensionTypeによって追加されたコンストラクターは、各フィールドを名前付き引数として受け取ります(クラス定義にリストされている順序で)。このコンストラクターは、各パラメーターの型チェックを行い、必要に応じてそれらを変換します。特に、 Tensorフィールドはtf.convert_to_tensorを使用して変換されます。 Tupleフィールドはtupleに変換されます。 Mappingフィールドは不変のdictに変換されます。

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)

フィールド値を宣言された型に変換できない場合、コンストラクターはTypeErrorを発生させます。

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

フィールドのデフォルト値は、その値をクラスレベルで設定することで指定できます。

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

印刷可能な表現

ExtensionTypeは、クラス名と各フィールドの値を含むデフォルトの印刷可能な表現メソッド( __repr__ )を追加します。

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

等式演算子

ExtensionTypeは、2つの値が同じタイプであり、すべてのフィールドが等しい場合に等しいと見なすデフォルトの等式演算子( __eq__および__ne__ )を追加します。テンソル場が同じ形状であり、すべての要素で要素ごとに等しい場合、テンソル場は等しいと見なされます。

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

検証方法

ExtensionTypeは、 __validate__メソッドを追加します。このメソッドは、フィールドの検証チェックを実行するためにオーバーライドできます。これは、コンストラクターが呼び出された後、およびフィールドが型チェックされて宣言された型に変換された後に実行されるため、すべてのフィールドに宣言された型があると見なすことができます。

次の例では、 dtypeを更新して、そのフィールドのshapeMaskedTensorを検証します。

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

強制不変性

ExtensionTypeは、__ setattr__メソッドと__delattr__メソッドをオーバーライドして、 __setattr__を防ぎ、拡張タイプの値が不変であることを保証します。

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

ExtensionTypeクラスには、対応するTypeSpecクラスがあり、これは自動的に作成され、 <extension_type_name>.Specとして保存されます。

このクラスは、ネストされたテンソルの値を除いて、値からすべての情報をキャプチャします。特に、値のTypeSpecは、ネストされたTensor、ExtensionType、またはCompositeTensorをその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)})
プレースホルダー29

TypeSpec値は明示的に作成することも、 tf.type_spec_from_valueを使用してExtensionType値から作成することもできます。

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

TypeSpecは、値を静的コンポーネント動的コンポーネントに分割するためにTensorFlowによって使用されます。

  • 静的コンポーネント(グラフ作成時に固定されます)は、 tf.TypeSpecでエンコードされます。
  • 動的コンポーネント(グラフが実行されるたびに変化する可能性があります)は、 tf.Tensorのリストとしてエンコードされます。

たとえば、 tf.functionは、引数に以前に表示されなかった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>}))

詳細については、 tf.functionガイドを参照してください。

ExtensionTypesのカスタマイズ

フィールドとそのタイプを単に宣言することに加えて、拡張タイプは次のことを行うことができます。

  • デフォルトの印刷可能な表現( __repr__ )をオーバーライドします。
  • メソッドを定義します。
  • classmethodsとstaticmethodsを定義します。
  • プロパティを定義します。
  • デフォルトのコンストラクター( __init__ )をオーバーライドします。
  • デフォルトの等式演算子( __eq__ )をオーバーライドします。
  • 演算子( __add____lt__など)を定義します。
  • フィールドのデフォルト値を宣言します。
  • サブクラスを定義します。

デフォルトの印刷可能な表現をオーバーライドする

拡張タイプのこのデフォルトの文字列変換演算子をオーバーライドできます。次の例では、 MaskedTensorクラスを更新して、値が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]]>

メソッドの定義

拡張タイプは、通常のPythonクラスと同じように、メソッドを定義できます。たとえば、 MaskedTensorタイプは、マスクされた値が指定されたdefault値に置き換えられたselfのコピーを返すwith_defaultメソッドを定義できます。オプションで、メソッドに@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)>

classmethodsとstaticmethodsの定義

拡張タイプは、 @staticmethod @classmethodを使用してメソッドを定義できます。たとえば、 MaskedTensorタイプは、指定された値で任意の要素をマスクするファクトリメソッドを定義できます。

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

プロパティの定義

拡張タイプは、通常のPythonクラスと同様に、 @propertyデコレータを使用してプロパティを定義できます。たとえば、 MaskedTensorタイプは、値のdtypeの省略形であるdtypeプロパティを定義できます。

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

デフォルトのコンストラクターをオーバーライドする

拡張タイプのデフォルトのコンストラクターをオーバーライドできます。カスタムコンストラクターは、宣言されたすべてのフィールドに値を設定する必要があります。カスタムコンストラクターが戻った後、すべてのフィールドが型チェックされ、値は上記のように変換されます。

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

または、デフォルトのコンストラクターをそのままにして、1つ以上のファクトリメソッドを追加することを検討することもできます。例えば:

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>)
プレースホルダー49

デフォルトの等式演算子のオーバーライド( __eq__

拡張タイプのデフォルトの__eq__演算子をオーバーライドできます。次の例では、 MaskedTensorを更新して、等しいかどうかを比較するときにマスクされた要素を無視します。

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)

前方参照の使用

フィールドのタイプがまだ定義されていない場合は、代わりにタイプの名前を含む文字列を使用できます。次の例では、 Nodeタイプがまだ(完全に)定義されていないため、文字列"Node"を使用してchildrenフィールドに注釈を付けています。

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

サブクラスの定義

拡張タイプは、標準のPython構文を使用してサブクラス化できます。拡張タイプのサブクラスは、新しいフィールド、メソッド、およびプロパティを追加する場合があります。コンストラクター、印刷可能な表現、および等価演算子をオーバーライドできます。次の例では、3つのTensorフィールドを使用してノード間のエッジのセットをエンコードする基本的なTensorGraphクラスを定義します。次に、各ノードの「特徴値」を記録するTensorフィールドを追加するサブクラスを定義します。サブクラスは、エッジに沿ってフィーチャ値を伝播するメソッドも定義します。

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)

プライベートフィールドの定義

拡張タイプのフィールドは、(標準のPython規則に従って)アンダースコアを前に付けることでプライベートとしてマークできます。これは、TensorFlowがフィールドを処理する方法にはまったく影響しません。ただし、これらのフィールドがプライベートであるという拡張タイプのユーザーへのシグナルとして機能するだけです。

ExtensionTypeのTypeSpecのカスタマイズ

ExtensionTypeクラスには、対応するTypeSpecクラスがあり、これは自動的に作成され、 <extension_type_name>.Specとして保存されます。詳細については、上記の「ネストされたTypeSpec」のセクションを参照してください。

TypeSpecをカスタマイズするには、 Specという名前の独自のネストされたクラスを定義するだけで、 ExtensionTypeはそれを自動的に構築されたTypeSpecの基礎として使用します。 Specクラスは次の方法でカスタマイズできます。

  • デフォルトの印刷可能な表現をオーバーライドします。
  • デフォルトのコンストラクターをオーバーライドします。
  • メソッド、classmethods、staticmethods、およびプロパティを定義します。

次の例では、 MaskedTensor.Specクラスをカスタマイズして、使いやすくします。

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)

TensorAPIディスパッチ

拡張タイプは、 tf.Tensorタイプによって定義されたインターフェースを特殊化または拡張するという意味で、「テンソルのような」ものにすることができます。テンソルのような拡張タイプの例には、 RaggedTensorSparseTensorMaskedTensorが含まれます。ディスパッチデコレータを使用して、テンソルのような拡張タイプに適用した場合のTensorFlow操作のデフォルトの動作をオーバーライドできます。 TensorFlowは現在、3つのディスパッチデコレータを定義しています。

単一のAPIのディスパッチ

tf.experimental.dispatch_for_apiデコレータは、指定されたシグニチャで呼び出されたときに、指定されたTensorFlow操作のデフォルトの動作をオーバーライドします。たとえば、このデコレータを使用して、 tf.stackMaskedTensor値を処理する方法を指定できます。

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

これは、 MaskedTensor値のリストで呼び出されるたびにtf.stackのデフォルトの実装をオーバーライドします( values引数には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]]>

tf.stackMaskedTensorTensorの混合値のリストを処理できるようにするには、 valuesパラメーターの型アノテーションを調整し、関数の本体を適切に更新します。

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

オーバーライドできるAPIのリストについては、 tf.experimental.dispatch_for_apiのAPIドキュメントを参照してください。

すべての単項要素ごとのAPIのディスパッチ

tf.experimental.dispatch_for_unary_elementwise_apisデコレータは、最初の引数(通常はxという名前)の値が型注釈x_typeと一致する場合は常に、すべての単項要素ごとのops( tf.math.cosなど)のデフォルトの動作をオーバーライドします。装飾された関数は2つの引数を取る必要があります。

  • api_func :単一のパラメーターを受け取り、要素ごとの操作を実行する関数( tf.abs )。
  • x :要素ごとの操作の最初の引数。

次の例では、 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)

この関数は、 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]>

バイナリのすべての要素ごとのAPIのディスパッチ

同様に、 tf.experimental.dispatch_for_binary_elementwise_apisを使用して、 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], [_, _, _]]>

オーバーライドされる要素ごとのAPIのリストについては、 tf.experimental.dispatch_for_unary_elementwise_apisおよびtf.experimental.dispatch_for_binary_elementwise_apisのAPIドキュメントを参照してください。

バッチ可能なExtensionTypes

単一のインスタンスを使用して値のバッチを表すことができる場合、 ExtensionTypeバッチ可能です。通常、これは、ネストされたすべてのTensorにバッチディメンションを追加することで実現されます。次のTensorFlowAPIでは、拡張タイプの入力がバッチ可能である必要があります。

  • tf.data.Datasetbatchunbatchfrom_tensor_slices
  • tf.Kerasfitevaluatepredict
  • tf.map_fn

デフォルトでは、 BatchableExtensionTypeは、ネストされたTensorCompositeTensor 、およびExtensionTypeをバッチ処理することによってバッチ値を作成します。これがクラスに適していない場合は、 tf.experimental.ExtensionTypeBatchEncoderを使用して、このデフォルトの動作をオーバーライドする必要があります。たとえば、個々のスパーステンソルのvaluesindicesdense_shapeフィールドを単純にスタックして、 tf.SparseTensor値のバッチを作成することは適切ではありません。ほとんどの場合、これらのテンソルは互換性のない形状であるため、スタックできません。 ;できたとしても、結果は有効なSparseTensorではありません。

BatchableExtensionTypeの例:ネットワーク

例として、負荷分散に使用される単純なNetworkクラスについて考えてみます。これは、各ノードで実行する必要のある作業量と、ノード間で作業を移動するために使用できる帯域幅を追跡します。

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

このタイプをバッチ可能にするには、基本タイプをBatchableExtensionTypeに変更し、オプションのバッチディメンションを含めるように各フィールドの形状を調整します。次の例では、バッチ形状を追跡するためのshapeフィールドも追加しています。このshapeフィールドは、 tf.data.Datasetまたはtf.map_fnには必要ありませんが、 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.]]]>
プレースホルダー73

次に、 tf.data.Datasetを使用して、ネットワークのバッチを反復処理できます。

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

また、 map_fnを使用して、各バッチ要素に関数を適用することもできます。

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

ExtensionTypesをサポートするTensorFlowAPI

@ tf.function

tf.functionは、Python関数のTensorFlowグラフを事前に計算するデコレータであり、TensorFlowコードのパフォーマンスを大幅に向上させることができます。拡張タイプの値は、 @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)>

tf.functioninput_signatureを明示的に指定する場合は、拡張タイプの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)>)

具体的な機能

具体的な関数は、 tf.functionによって作成された個々のトレースされたグラフをカプセル化します。拡張タイプは、具象関数で透過的に使用できます。

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

制御フロー操作

拡張タイプは、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, _]>

サインコントロールフロー

拡張タイプは、tf.functionの制御フローステートメントでもサポートされています(サインを使用)。次の例では、 ifステートメントとforステートメントは、拡張タイプをサポートするtf.condおよびtf.while_loop操作に自動的に変換されます。

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

ケラス

tf.kerasは、ディープラーニングモデルを構築およびトレーニングするためのTensorFlowの高レベルAPIです。拡張タイプは、Kerasモデルへの入力として渡され、Kerasレイヤー間で渡され、Kerasモデルによって返される場合があります。 Kerasは現在、拡張タイプに2つの要件を課しています。

  • それらはバッチ可能でなければなりません(上記の「バッチ可能なExtensionTypes」を参照)。
  • には、 shapeという名前のフィールドまたはプロパティが必要です。 shape[0]はバッチディメンションと見なされます。

次の2つのサブセクションでは、拡張タイプをKerasで使用する方法を示す例を示します。

Kerasの例: Network

最初の例では、上記の「Batchable ExtensionTypes」セクションで定義されているNetworkクラスについて考えてみます。これは、ノード間の負荷分散作業に使用できます。その定義はここで繰り返されます:

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

Networkを処理する新しいKerasレイヤーを定義できます。

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)

次に、このレイヤーを使用して単純なモデルを作成できます。 ExtensionTypeをモデルにフィードするには、 TypeSpec type_spec設定されたtf.keras.layer.Inputレイヤーを使用できます。 type_specモデルを使用してバッチを処理する場合は、type_specにバッチディメンションを含める必要があります。

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

最後に、モデルを単一のネットワークとネットワークのバッチに適用できます。

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

Kerasの例:MaskedTensor

この例では、 MaskedTensorKerasをサポートするように拡張されています。 shapeは、 valuesフィールドから計算されるプロパティとして定義されます。 Kerasでは、このプロパティを拡張タイプとそのTypeSpecの両方に追加する必要があります。 MaskedTensor__name__変数も定義します。これは、 SavedModelのシリアル化(以下)に必要です。

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

次に、ディスパッチデコレータを使用して、いくつかのTensorFlowAPIのデフォルトの動作をオーバーライドします。これらのAPIは標準のKerasレイヤー( Denseレイヤーなど)で使用されるため、これらをオーバーライドすると、 MaskedTensorでこれらのレイヤーを使用できるようになります。この例の目的のために、マスクされたテンソルのmatmulは、マスクされた値をゼロとして扱うように定義されます(つまり、それらを積に含めないようにします)。

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

次に、標準のKerasレイヤーを使用して、MaskedTensor入力を受け入れるMaskedTensorモデルを構築できます。

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

SavedModelは、重みと計算の両方を含む、シリアル化されたTensorFlowプログラムです。 Kerasモデルまたはカスタムモデルから構築できます。いずれの場合も、拡張タイプは、SavedModelで定義された関数とメソッドで透過的に使用できます。

SavedModelは、拡張タイプに__name__フィールドがある限り、拡張タイプを処理するモデル、レイヤー、および関数を保存できます。この名前は拡張タイプを登録するために使用されるため、モデルのロード時に見つけることができます。

例:Kerasモデルの保存

拡張タイプを使用するKerasモデルは、 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)>

例:カスタムモデルの保存

SavedModelを使用して、拡張タイプを処理する関数を含むカスタムtf.Moduleサブクラスを保存することもできます。

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

ExtensionTypeが使用できない場合のSavedModelのロード

ExtensionTypeを使用するSavedModelをロードしたが、そのExtensionTypeが使用できない(つまり、インポートされていない)場合、警告が表示され、TensorFlowは「匿名拡張タイプ」オブジェクトの使用にフォールバックします。このオブジェクトには元のタイプと同じフィールドがありますが、カスタムメソッドやプロパティなど、タイプに追加したカスタマイズはありません。

TensorFlowサービングでExtensionTypesを使用する

現在、 TensorFlowサービング(およびSavedModel「signatures」ディクショナリの他のコンシューマー)では、すべての入力と出力が生のテンソルである必要があります。拡張タイプを使用するモデルでTensorFlowサービングを使用する場合は、テンソルから拡張タイプ値を作成または分解するラッパーメソッドを追加できます。例えば:

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>

データセット

tf.dataは、単純で再利用可能な部分から複雑な入力パイプラインを構築できるようにするAPIです。そのコアデータ構造はtf.data.Datasetであり、これは要素のシーケンスを表し、各要素は1つ以上のコンポーネントで構成されます。

拡張タイプを使用したデータセットの構築

データセットは、 Dataset.from_tensorsDataset.from_tensor_slices 、または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]>

拡張タイプを使用したデータセットのバッチ処理とバッチ解除

拡張タイプのデータセットは、Dataset.batchおよびDataset.unbatchを使用してDataset.batch処理および非バッチ処理できます。

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