확장 유형

TensorFlow.org에서 보기 Google Colab에서 실행 GitHub에서 소스 보기 노트북 다운로드

설정

!pip install -q tf_nightly
import tensorflow as tf
import numpy as np
from typing import Tuple, List, Mapping, Union, Optional
import tempfile
2022-12-14 22:14:29.901620: E tensorflow/tsl/lib/monitoring/collection_registry.cc:81] Cannot register 2 metrics with the same name: /tensorflow/core/bfc_allocator_delay

확장 유형

사용자 정의 유형을 사용하면 프로젝트를 더 읽기 쉽고 모듈식으로 유지 관리할 수 있습니다. 그러나 대부분의 TensorFlow API는 사용자 정의 Python 유형에 대한 지원이 매우 제한적입니다. 이것은 (예 모두 높은 수준의 API를 포함 Keras , tf.function , tf.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__

일반적으로 확장 유형은 다음 두 가지 범주 중 하나로 분류되는 경향이 있습니다.

  • 관련 값의 컬렉션을 그룹화하고 해당 값을 기반으로 유용한 작업을 제공할 수 있는 데이터 구조. 데이터 구조는 상당히 일반적일 수 있습니다(예 TensorGraph 예). 또는 특정 모델에 고도로 맞춤화될 수 있습니다.

  • "Tensor"의 개념을 전문화하거나 확장하는 Tensor와 유사한 유형입니다. 이 범주의 유형에는 rank , shape 및 일반적으로 dtype . tf.stack , tf.add 또는 tf.matmul )과 함께 사용하는 것이 좋습니다. MaskedTensorCSRSparseMatrix 는 텐서 유사 유형의 예입니다.

지원되는 API

확장 유형은 다음 TensorFlow API에서 지원됩니다.

  • Keras ModelsLayers 대한 입력 및 출력으로 사용할 수 있습니다.
  • tf.data.Dataset : 확장 유형은 Datasets Iterators 의해 반환됩니다.
  • TensorFlow 허브: 확장 유형을 tf.hub 모듈의 입력 및 출력으로 사용할 수 있습니다.
  • SavedModel SavedModel 함수에 대한 입력 및 출력으로 사용할 수 있습니다.
  • tf.function @tf.function 데코레이터로 래핑된 함수의 인수 및 반환 값으로 사용할 수 있습니다.
  • While 루프: 확장 유형을 tf.while_loop에서 루프 변수로 사용할 수 있으며, while 루프 본문에서 인수 및 반환 값으로도 사용할 수 있습니다.
  • 조건: tf.condtf.case를 사용하여 확장 유형을 조건부로 선택할 수 있습니다.
  • tf.py_function: 확장 유형을 인수로 사용하고 func 인수의 값을 tf.py_function에 반환할 수 있습니다.
  • Tensor 연산: 텐서 입력을 허용하는 대부분의 TensorFlow 연산(예: tf.matmul, tf.gather, tf.reduce_sum)을 지원하도록 확장 유형을 확장할 수 있습니다. 자세한 정보는 아래의 "디스패치" 섹션으로 이동하여 확인할 수 있습니다.
  • 분산 전략: 확장 유형을 복제본별 값으로 사용할 수 있습니다.

자세한 내용은 아래 "ExtensionTypes를 지원하는 TensorFlow API" 섹션을 참조하세요.

요구 사항

필드 유형

모든 필드(인스턴스 변수)를 선언해야 하며 각 필드에 유형 주석을 제공해야 합니다. 지원되는 유형 주석은 다음과 같습니다.

유형 예시
파이썬 정수 i: int
파이썬 수레 f: float
파이썬 문자열 s: str
파이썬 부울 b: bool
파이썬 없음 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]
가변 길이 튜플 lengths: typing.Tuple[int, ...]
매핑 tags: typing.Mapping[str, tf.Tensor]
선택적 값 weight: typing.Optional[tf.Tensor]

가변성

확장 유형은 변경 불가능해야 합니다. 이렇게 하면 TensorFlow의 그래프 추적 메커니즘으로 적절하게 추적할 수 있습니다. 확장 유형 값을 변경하려는 경우 값을 변환하는 메서드를 대신 정의하는 것이 좋습니다. 예를 들어 MaskedTensor 를 변경하기 위해 set_mask MaskedTensor 를 반환하는 replace_mask 메서드를 정의할 수 있습니다.

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 .
  • 텐서 API 디스패치 지원.

이 기능을 사용자 정의하는 방법에 대한 자세한 정보는 아래의 'ExtensionType 사용자 정의하기' 섹션으로 이동하여 확인합니다.

건설자

ExtensionType 에 의해 추가된 생성자는 각 필드를 명명된 인수로 사용합니다(클래스 정의에 나열된 순서대로). 이 생성자는 각 매개변수를 유형 검사하고 필요한 경우 변환합니다. 특히, Tensor tf.convert_to_tensor 사용하여 변환됩니다. Tuple 필드로 변환됩니다 tuple 의; Mapping 필드는 변경할 수 없는 사전으로 변환됩니다.

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.
# For example, `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 'NoneType'

필드의 기본값은 클래스 수준에서 값을 설정하여 지정할 수 있습니다.

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 은 유형이 동일하고 모든 필드가 동일한 경우 두 값을 동일하게 간주하는 기본 동등 연산자( __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

참고: Tensor 가 포함된 경우 __eq__ 는 (Python 부울 값 대신) Tensor 반환할 수 있습니다.

검증 방법

ExtensionType은 필드에 대한 유효성 검사를 수행하기 위해 재정의할 수 있는 __validate__ 메서드를 추가합니다. 이는 생성자가 호출되고 필드의 유형을 검사하고 선언된 유형으로 변환된 후에 실행되므로 모든 필드에 선언된 유형이 있다고 가정할 수 있습니다.

다음 예제는 MaskedTensor를 업데이트하여 해당 필드의 shapedtype의 유효성을 검사합니다.

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__ 메서드를 재정의하여 변형을 방지하여 확장 유형 값을 변경할 수 없도록 합니다.

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.

중첩된 유형 사양

ExtensionType 클래스에는 자동으로 생성되고 <extension_type_name>.Spec TypeSpec 클래스가 있습니다.

이 클래스는 중첩된 텐서의 값을 제외한 값에서 모든 정보를 캡처합니다. 특히 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.
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 값은 명시적으로 구성하거나 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 .

@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}))
<<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 가이드를 참조하십시오.

ExtensionType 사용자 정의

단순히 필드와 해당 유형을 선언하는 것 외에도 확장 유형은 다음을 수행할 수 있습니다.

  • 기본 인쇄 가능한 표현( __repr__ )을 재정의합니다.
  • 방법을 정의합니다.
  • 클래스 메서드와 정적 메서드를 정의합니다.
  • 속성을 정의합니다.
  • 기본 생성자( __init__ )를 재정의합니다.
  • 기본 항등 연산자( __eq__ )를 재정의합니다.
  • 연산자를 정의합니다(예: __add____lt__ ).
  • 필드의 기본값을 선언합니다.
  • 하위 클래스를 정의합니다.

기본 인쇄 가능한 표현 재정의

확장 유형에 대해 이 기본 문자열 변환 연산자를 재정의할 수 있습니다. 다음 예제에서는 값이 Eager 모드에서 인쇄될 때 더 읽기 쉬운 문자열 표현을 생성 MaskedTensor

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

클래스 메서드 및 정적 메서드 정의

@classmethod@staticmethod 데코레이터를 사용하여 메소드를 정의할 수 있습니다. 예를 들어 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 [[1, _, 2], [3, _, _]]>

속성 정의

확장 유형은 일반 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>)

또는 기본 생성자를 그대로 두고 하나 이상의 팩토리 메서드를 추가하는 방법을 고려할 수 있습니다. 예제는 다음과 같습니다.

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

기본 항등 연산자 재정의( __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)

참고: 기본 구현은 단순히 __eq__ 를 호출하고 결과를 무효화하기 __ne__ 를 재정의할 필요가 없습니다.

정방향 참조 사용

필드 유형이 아직 정의되지 않은 경우 유형 이름이 포함된 문자열을 대신 사용할 수 있습니다. 다음 예제에서는 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 클래스에는 자동으로 생성되고 <extension_type_name>.Spec TypeSpec 클래스가 있습니다. 자세한 내용은 위의 "중첩된 TypeSpec" 섹션을 참조하세요.

TypeSpec 을 사용자 정의하려면 Spec 이라는 자체 중첩 클래스를 정의하기만 하면 ExtensionType 이 이를 자동으로 생성된 TypeSpec 의 기초로 사용합니다. Spec 클래스를 사용자 정의할 수 있습니다.

  • 기본 인쇄 가능한 표현을 재정의합니다.
  • 기본 생성자를 재정의합니다.
  • 메서드, 클래스 메서드, 정적 메서드 및 속성을 정의합니다.

다음 예제에서는 사용하기 쉽도록 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)

참고 : 사용자 정의 Spec ExtensionType 선언되지 않은 인스턴스 변수를 사용할 수 없습니다.

텐서 API 디스패치

tf.Tensor 유형에 의해 정의된 인터페이스를 전문화하거나 확장한다는 점에서 "텐서와 유사"할 수 있습니다. 텐서와 유사한 확장 유형의 예로는 RaggedTensor , SparseTensorMaskedTensor 있습니다. 디스패치 데코레이터 는 텐서와 유사한 확장 유형에 적용될 때 TensorFlow 작업의 기본 동작을 재정의하는 데 사용할 수 있습니다. TensorFlow는 현재 세 가지 디스패치 데코레이터를 정의합니다.

단일 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.stack 이 혼합된 MaskedTensorTensor 값 목록을 처리할 수 있도록 하려면 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 과 일치할 때마다 모든 단항 요소별 연산(예: tf.math.cos )의 기본 동작을 재정의합니다. 데코레이팅된 함수는 두 개의 인수를 취해야 합니다.

  • 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_apistf.experimental.dispatch_for_binary_elementwise_apis에 대한 API 문서로 이동합니다.

일괄 처리 가능한 확장 유형

ExtensionType 단일 인스턴스 값의 배치를 나타내는 데 사용할 수있는 경우 batchable이다. Tensor 배치 차원을 추가하여 수행됩니다. 다음 TensorFlow API를 사용하려면 모든 확장 유형 입력이 일괄 처리 가능해야 합니다.

기본적으로 BatchableExtensionType Tensor , CompositeTensorExtensionType 일괄 처리하여 일괄 처리된 값을 생성합니다. 이것이 클래스에 적합하지 않은 경우 tf.experimental.ExtensionTypeBatchEncoder 를 사용하여 이 기본 동작을 재정의해야 합니다. 예를 들어, 개별 희소 텐서의 values , indicesdense_shape tf.SparseTensor 값의 배치를 만드는 것은 적절하지 않습니다. 대부분의 경우 이러한 텐서는 호환되지 않는 모양을 가지고 있기 때문에 스택할 수 없습니다. ; 가능하더라도 결과는 유효한 SparseTensor .

참고 : BatchableExtensionType tf.stack , tf.concat , tf.slice 등에 대한 디스패처를 자동으로 정의하지 않습니다 . 이러한 API에서 클래스를 지원해야 하는 경우 위에서 설명한 디스패치 데코레이터를 사용하세요.

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

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를 지원하는 TensorFlow API

@tf.function

tf.function 은 TensorFlow 코드의 성능을 크게 향상시킬 수 있는 Python 함수용 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)>

input_signature 대해 tf.function 를 명시적으로 지정 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의 제어 흐름 문에서도 지원됩니다(autograph 사용). 다음 예에서 if 문과 for 문은 확장 유형을 지원 tf.condtf.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는 현재 확장 유형에 두 가지 요구 사항을 적용합니다.

  • 배치 처리할 수 있어야 합니다(위의 "배치 처리할 수 있는 ExtensionType"으로 이동).
  • shape이라는 필드 또는 속성이 있어야 합니다. shape[0]은 배치 차원으로 간주됩니다.

다음 두 하위 섹션에서는 확장 유형을 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 with 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 the "Batchable `ExtensionType`s" section.
    return balance_work_greedy(inputs)

그런 다음 이러한 레이어를 사용하여 간단한 모델을 생성할 수 있습니다. type_spec가 있는 tf.keras.layer.Input를 확장 유형의 TypeSpec으로 설정할 수 있습니다. Keras 모델을 사용하여 배치 처리하는 경우 배치 차원에 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.]]]>

케라스 예시: MaskedTensor

이 예에서 MaskedTensor Keras 를 지원하도록 확장되었습니다. shape values 필드에서 계산되는 속성으로 정의됩니다. TypeSpec 모두에 이 속성을 추가해야 합니다. MaskedTensor SavedModel 직렬화에 필요한 __name__ 변수도 정의합니다(아래 참조).

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

다음으로 디스패치 데코레이터는 여러 TensorFlow API의 기본 동작을 재정의하는 데 사용합니다. 이러한 API는 표준 Keras 레이어(예: Dense 레이어)에서 사용하므로 이를 재정의하면 MaskedTensor와 함께 해당 레이어를 사용할 수 있습니다. 이 예제에서는 마스킹된 텐서의 matmul이 마스킹된 값을 0으로 처리하도록(즉, 제품에 포함하지 않도록) 정의합니다.

@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 입력을 허용하는 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 1s/step - loss: 6.4280
Epoch 2/3
1/1 [==============================] - 0s 5ms/step - loss: 6.1971
Epoch 3/3
1/1 [==============================] - 0s 5ms/step - loss: 6.1420
tf.Tensor(
[[-0.31530094]
 [ 0.79135597]
 [ 0.27101254]], shape=(3, 1), dtype=float32)

저장된 모델

SavedModel 은 가중치와 계산을 모두 포함하는 직렬화된 TensorFlow 프로그램입니다. Keras 모델 또는 사용자 지정 모델에서 구축할 수 있습니다. 두 경우 모두 확장 유형은 SavedModel에 의해 정의된 함수 및 메소드와 함께 투명하게 사용될 수 있습니다.

__name__ 필드가 있는 한 확장 유형을 처리하는 모델, 계층 및 함수를 저장할 수 있습니다. 이 이름은 확장 유형을 등록하는 데 사용되므로 모델을 로드할 때 찾을 수 있습니다.

예: 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)
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: /tmpfs/tmp/tmpt3nnr57a/assets
INFO:tensorflow:Assets written to: /tmpfs/tmp/tmpt3nnr57a/assets
<tf.Tensor: shape=(3, 1), dtype=float32, numpy=
array([[-0.31530094],
       [ 0.79135597],
       [ 0.27101254]], 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: /tmpfs/tmp/tmpw6f0yvf0/assets
INFO:tensorflow:Assets written to: /tmpfs/tmp/tmpw6f0yvf0/assets
<MaskedTensor [_, 200.0, _]>

ExtensionType을 사용할 수 없을 때 저장된 모델 로드

ExtensionType을 사용하는 SavedModel을 로드하였으나 해당 ExtensionType을 사용할 수 없는 경우(즉, 가져오지 않은 경우) 경고가 표시되고 TensorFlow는 "익명 확장 유형" 객체를 사용하도록 대체됩니다. 이 객체는 원래 유형과 동일한 필드를 갖지만 사용자 정의 메서드 또는 속성과 같은 유형에 추가한 추가 사용자 정의 설정이 부족합니다.

TensorFlow Serving으로 ExtensionType 사용하기

현재 TensorFlow Serving(및 SavedModel "서명" 사전의 다른 소비자)에서는 모든 입력과 출력이 원시 텐서여야 합니다. 확장 유형을 사용하는 모델과 함께 TensorFlow Serving을 사용하려는 경우 텐서로부터 확장 유형 값을 구성하거나 분해하는 래핑 메서드를 추가할 수 있습니다. 예제는 다음과 같습니다.

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: /tmpfs/tmp/tmp53zdnf3c/assets
INFO:tensorflow:Assets written to: /tmpfs/tmp/tmp53zdnf3c/assets
<tf.Tensor: shape=(), dtype=float32, numpy=12.0>

데이터세트

tf.data 는 간단하고 재사용 가능한 부분으로 복잡한 입력 파이프라인을 구축할 수 있는 API입니다. 핵심 데이터 구조는 tf.data.Dataset 이며, 이는 각 요소가 하나 이상의 구성 요소로 구성된 일련의 요소를 나타냅니다.

확장 유형으로 데이터세트 빌드

Dataset.from_tensors , Dataset.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.batchDataset.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]>