Các loại tiện ích mở rộng

Xem trên TensorFlow.org Chạy trong Google Colab Xem nguồn trên GitHubTải xuống sổ ghi chép

Thành lập

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

Các loại tiện ích mở rộng

Các kiểu do người dùng xác định có thể làm cho các dự án dễ đọc hơn, dễ đọc hơn, dễ bảo trì hơn. Tuy nhiên, hầu hết các API TensorFlow đều hỗ trợ rất hạn chế cho các loại Python do người dùng xác định. Điều này bao gồm cả các API cấp cao (chẳng hạn như Keras , tf. Chức năng, tf.SavedModel ) và các API cấp thấp hơn (chẳng hạn như tf. tf.while_loop và tf.concat ). Các kiểu mở rộng TensorFlow có thể được sử dụng để tạo các kiểu hướng đối tượng do người dùng xác định hoạt động trơn tru với các API của TensorFlow. Để tạo một kiểu mở rộng, chỉ cần xác định một lớp Python với tf.experimental.ExtensionType làm cơ sở của nó và sử dụng chú thích kiểu để chỉ định kiểu cho từng trường.

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

Lớp cơ sở tf.experimental.ExtensionType hoạt động tương tự như typing.NamedTuple và @dataclasses.dataclass từ thư viện Python chuẩn. Đặc biệt, nó tự động thêm một phương thức khởi tạo và các phương thức đặc biệt (chẳng hạn như __repr__ và __eq__ ) dựa trên các chú thích kiểu trường.

Thông thường, các loại tiện ích mở rộng có xu hướng thuộc một trong hai loại:

  • Cấu trúc dữ liệu , nhóm các giá trị có liên quan lại với nhau và có thể cung cấp các hoạt động hữu ích dựa trên các giá trị đó. Cấu trúc dữ liệu có thể khá chung chung (chẳng hạn như ví dụ TensorGraph ở trên); hoặc chúng có thể được tùy chỉnh cao cho một mô hình cụ thể.

  • Loại giống như kéo căng , chuyên biệt hóa hoặc mở rộng khái niệm "Tensor". Các loại trong danh mục này có một rank , một shape và thường là một dtype ; và sử dụng chúng với các hoạt động Tensor (chẳng hạn như tf.stack , tf.add hoặc tf.matmul ). MaskedTensor và CSRSparseMatrix là những ví dụ về các loại giống tensor.

Các API được hỗ trợ

Các loại tiện ích mở rộng được hỗ trợ bởi các API TensorFlow sau:

  • Keras : Các loại phần mở rộng có thể được sử dụng làm đầu vào và đầu ra cho Models và Layers Keras.
  • tf.data.Dataset : Các kiểu mở rộng có thể được bao gồm trong Datasets và được trả về bởi Iterators tập dữ liệu.
  • Trung tâm Tensorflow : Các loại mở rộng có thể được sử dụng làm đầu vào và đầu ra cho mô-đun tf.hub .
  • SavedModel : Các loại tiện ích mở rộng có thể được sử dụng làm đầu vào và đầu ra cho các chức năng của SavedModel .
  • tf . functions: Các loại phần mở rộng có thể được sử dụng làm đối số và trả về giá trị cho các hàm được bao bọc bằng trình trang trí @tf.function . function.
  • vòng lặp while : Các kiểu mở rộng có thể được sử dụng làm biến vòng lặp trong tf.while_loop và có thể được sử dụng làm đối số và trả về giá trị cho phần thân của vòng lặp while.
  • điều kiện : Các loại phần mở rộng có thể được chọn có điều kiện bằng cách sử dụng tf.cond và tf.case .
  • py_ Chức năng : Các loại phần mở rộng có thể được sử dụng làm đối số và trả về giá trị cho đối số func cho tf.py_function .
  • Hoạt động kéo dài: Các loại tiện ích mở rộng có thể được mở rộng để hỗ trợ hầu hết các hoạt động TensorFlow chấp nhận đầu vào Tensor (ví dụ: tf.matmul , tf.gather và tf.reduce_sum ). Xem phần " Công văn " bên dưới để biết thêm thông tin.
  • chiến lược phân phối : Các loại tiện ích mở rộng có thể được sử dụng làm giá trị cho mỗi bản sao.

Để biết thêm chi tiết, hãy xem phần "TensorFlow APIs hỗ trợ ExtensionTypes" bên dưới.

Yêu cầu

Các loại trường

Tất cả các trường (hay còn gọi là biến phiên bản) phải được khai báo và phải cung cấp chú thích kiểu cho mỗi trường. Các loại chú thích sau được hỗ trợ:

Loại Thí dụ
Số nguyên trong Python i: int
Trăn nổi f: float
Chuỗi Python s: str
Các boolean trong Python b: bool
Python Không có n: None
Tensor hình dạng shape: tf.TensorShape
Tensor dtypes dtype: tf.DType
Căng thẳng t: tf.Tensor
Các loại tiện ích mở rộng mt: MyMaskedTensor
Ragged Tensors rt: tf.RaggedTensor
Độ căng thưa st: tf.SparseTensor
Lát được lập chỉ mục s: tf.IndexedSlices
Độ căng tùy chọn o: tf.experimental.Optional
Nhập công đoàn int_or_float: typing.Union[int, float]
Tuples params: typing.Tuple[int, float, tf.Tensor, int]
Tuples chiều dài biến đổi lengths: typing.Tuple[int, ...]
Ánh xạ tags: typing.Mapping[str, tf.Tensor]
Giá trị tùy chọn weight: typing.Optional[tf.Tensor]

Tính đột biến

Các loại tiện ích mở rộng bắt buộc phải là bất biến. Điều này đảm bảo rằng chúng có thể được theo dõi đúng cách bằng cơ chế truy tìm đồ thị của TensorFlow. Nếu bạn thấy mình muốn thay đổi giá trị kiểu tiện ích, thay vào đó hãy xem xét xác định các phương pháp biến đổi giá trị. Ví dụ: thay vì xác định phương thức set_mask để thay đổi MaskedTensor , bạn có thể xác định phương thức replace_mask trả về một MaskedTensor mới:

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)

Chức năng được thêm bởi ExtensionType

Lớp cơ sở ExtensionType cung cấp chức năng sau:

  • Một hàm tạo ( __init__ ).
  • Phương thức biểu diễn có thể in được ( __repr__ ).
  • Toán tử bình đẳng và bất bình đẳng ( __eq__ ).
  • Phương thức xác thực ( __validate__ ).
  • Tính bất biến bắt buộc.
  • Một TypeSpec lồng nhau.
  • Hỗ trợ điều phối API Tensor.

Xem phần "Tùy chỉnh Loại mở rộng" bên dưới để biết thêm thông tin về cách tùy chỉnh chức năng này.

Constructor

Hàm tạo được thêm bởi ExtensionType nhận mỗi trường làm đối số được đặt tên (theo thứ tự chúng được liệt kê trong định nghĩa lớp). Hàm tạo này sẽ kiểm tra từng tham số và chuyển đổi chúng khi cần thiết. Đặc biệt, các trường Tensor được chuyển đổi bằng cách sử dụng tf.convert_to_tensor ; Các trường Tuple được chuyển đổi thành tuple s; và các trường Mapping được chuyển đổi thành các ô bất biến.

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)

Hàm tạo tăng TypeError nếu một giá trị trường không thể được chuyển đổi thành kiểu đã khai báo của nó:

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

Giá trị mặc định cho một trường có thể được chỉ định bằng cách đặt giá trị của nó ở cấp lớp:

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

Bản trình bày có thể in được

ExtensionType thêm một phương thức biểu diễn có thể in mặc định ( __repr__ ) bao gồm tên lớp và giá trị cho mỗi trường:

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

Toán tử bình đẳng

ExtensionType thêm các toán tử bình đẳng mặc định ( __eq__ và __ne__ ) coi hai giá trị bằng nhau nếu chúng có cùng kiểu và tất cả các trường của chúng đều bằng nhau. Các trường kéo căng được coi là bằng nhau nếu chúng có cùng hình dạng và bằng nhau theo từng phần tử đối với tất cả các phần tử.

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

Phương pháp xác thực

ExtensionType thêm một phương thức __validate__ , phương thức này có thể được ghi đè để thực hiện kiểm tra xác thực trên các trường. Nó được chạy sau khi phương thức khởi tạo được gọi, và sau khi các trường đã được kiểm tra kiểu và chuyển đổi thành kiểu đã khai báo của chúng, vì vậy nó có thể giả định rằng tất cả các trường đều có kiểu được khai báo.

ví dụ sau anh ấy cập nhật MaskedTensor để xác thực shape và dtype của các trường của nó:

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

Tính bất biến bắt buộc

ExtensionType ghi đè các phương thức __setattr__ và __delattr__ để ngăn đột biến, đảm bảo rằng các giá trị kiểu mở rộng là không thay đổi.

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.

Loại lồng nhauSpec

Mỗi lớp ExtensionType có một lớp TypeSpec tương ứng, lớp này được tạo tự động và được lưu trữ dưới dạng <extension_type_name>.Spec .

Lớp này nắm bắt tất cả thông tin từ một giá trị ngoại trừ các giá trị của bất kỳ tenxơ lồng nhau nào. Cụ thể, TypeSpec cho một giá trị được tạo bằng cách thay thế bất kỳ Tensor, ExtensionType hoặc CompositeTensor lồng nhau nào bằng TypeSpec của nó.

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

Giá trị TypeSpec có thể được tạo một cách rõ ràng hoặc chúng có thể được tạo từ giá trị ExtensionType bằng cách sử dụng tf.type_spec_from_value :

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

TypeSpec được TensorFlow sử dụng để phân chia các giá trị thành một thành phần tĩnh và một thành phần động :

  • Thành phần tĩnh (được cố định tại thời điểm xây dựng đồ thị) được mã hóa bằng tf.TypeSpec .
  • Thành phần động (có thể thay đổi mỗi khi chạy đồ thị) được mã hóa dưới dạng danh sách các tf.Tensor s.

Ví dụ: tf.function . function khôi phục lại hàm được bao bọc của nó bất cứ khi nào một đối số có TypeSpec chưa từng thấy trước đó:

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

Để biết thêm thông tin, hãy xem Hướng dẫn chức năng tf .

Tùy chỉnh ExtensionTypes

Ngoài việc chỉ khai báo các trường và kiểu của chúng, các kiểu mở rộng có thể:

  • Ghi đè biểu diễn có thể in mặc định ( __repr__ ).
  • Xác định các phương pháp.
  • Xác định classmethods và staticmethods.
  • Xác định thuộc tính.
  • Ghi đè hàm tạo mặc định ( __init__ ).
  • Ghi đè toán tử bình đẳng mặc định ( __eq__ ).
  • Xác định các toán tử (chẳng hạn như __add__ và __lt__ ).
  • Khai báo giá trị mặc định cho các trường.
  • Xác định các lớp con.

Ghi đè biểu diễn có thể in mặc định

Bạn có thể ghi đè toán tử chuyển đổi chuỗi mặc định này cho các loại tiện ích mở rộng. Ví dụ sau cập nhật lớp MaskedTensor để tạo biểu diễn chuỗi dễ đọc hơn khi các giá trị được in ở chế độ 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]]>

Xác định phương pháp

Các kiểu mở rộng có thể xác định các phương thức, giống như bất kỳ lớp Python bình thường nào. Ví dụ: kiểu MaskedTensor có thể xác định một phương thức with_default trả về một bản sao của self nó với các giá trị bị che được thay thế bằng một giá trị default nhất định. Các phương pháp có thể được chú thích tùy ý với trình trang trí chức @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)>

Định nghĩa classmethods và staticmethods

Các loại tiện ích mở rộng có thể xác định các phương thức bằng cách sử dụng trình trang trí @classmethod và @staticmethod . Ví dụ: kiểu MaskedTensor có thể xác định một phương thức gốc che bất kỳ phần tử nào có giá trị nhất định:

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

Xác định thuộc tính

Các loại tiện ích mở rộng có thể xác định các thuộc tính bằng cách sử dụng trình trang trí @property , giống như bất kỳ lớp Python bình thường nào. Ví dụ: kiểu MaskedTensor có thể xác định thuộc tính dtype là cách viết tắt của loại giá trị:

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

Ghi đè hàm tạo mặc định

Bạn có thể ghi đè hàm tạo mặc định cho các loại tiện ích mở rộng. Các hàm tạo tùy chỉnh phải đặt một giá trị cho mọi trường được khai báo; và sau khi phương thức khởi tạo tùy chỉnh trả về, tất cả các trường sẽ được kiểm tra kiểu và các giá trị sẽ được chuyển đổi như mô tả ở trên.

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

Ngoài ra, bạn có thể cân nhắc giữ nguyên phương thức khởi tạo mặc định nhưng thêm một hoặc nhiều phương thức gốc. Ví dụ:

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

Ghi đè toán tử bình đẳng mặc định ( __eq__ )

Bạn có thể ghi đè toán tử __eq__ mặc định cho các loại tiện ích mở rộng. Ví dụ sau cập nhật MaskedTensor để bỏ qua các phần tử bị che khi so sánh về sự bình đẳng.

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)

Sử dụng tham chiếu chuyển tiếp

Nếu loại cho một trường chưa được xác định, bạn có thể sử dụng một chuỗi chứa tên của loại thay thế. Trong ví dụ sau, chuỗi "Node" được sử dụng để chú thích trường children vì loại Node chưa được xác định (đầy đủ).

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

Định nghĩa các lớp con

Các loại phần mở rộng có thể được phân lớp bằng cách sử dụng cú pháp Python chuẩn. Các lớp con kiểu mở rộng có thể thêm các trường, phương thức và thuộc tính mới; và có thể ghi đè hàm tạo, biểu diễn có thể in và toán tử bình đẳng. Ví dụ sau định nghĩa một lớp TensorGraph cơ bản sử dụng ba trường Tensor để mã hóa một tập hợp các cạnh giữa các nút. Sau đó, nó định nghĩa một lớp con bổ sung trường Tensor để ghi lại "giá trị tính năng" cho mỗi nút. Lớp con cũng định nghĩa một phương thức để truyền các giá trị đặc trưng dọc theo các cạnh.

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)

Xác định các trường riêng tư

Các trường của một loại tiện ích mở rộng có thể được đánh dấu là riêng tư bằng cách thêm dấu gạch dưới vào đầu chúng (tuân theo các quy ước chuẩn của Python). Điều này không ảnh hưởng đến cách TensorFlow xử lý các trường theo bất kỳ cách nào; mà chỉ đóng vai trò như một tín hiệu cho bất kỳ người dùng loại tiện ích mở rộng nào rằng các trường đó là riêng tư.

Tùy chỉnh TypeSpec của ExtensionType

Mỗi lớp ExtensionType có một lớp TypeSpec tương ứng, lớp này được tạo tự động và được lưu trữ dưới dạng <extension_type_name>.Spec . Để biết thêm thông tin, hãy xem phần "TypeSpec lồng nhau" ở trên.

Để tùy chỉnh TypeSpec , chỉ cần xác định lớp lồng nhau của riêng bạn có tên là Spec và ExtensionType sẽ sử dụng nó làm cơ sở cho TypeSpec xây dựng tự động. Bạn có thể tùy chỉnh lớp Spec bằng cách:

  • Ghi đè biểu diễn có thể in mặc định.
  • Ghi đè hàm tạo mặc định.
  • Định nghĩa phương thức, classmethods, staticmethods và thuộc tính.

Ví dụ sau đây tùy chỉnh lớp MaskedTensor.Spec để làm cho nó dễ sử dụng hơn:

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)

Công văn Tensor API

Các kiểu mở rộng có thể "giống tensor", theo nghĩa là chúng chuyên biệt hóa hoặc mở rộng giao diện được xác định bởi kiểu tf.Tensor . Ví dụ về các loại phần mở rộng giống tensor bao gồm RaggedTensor , SparseTensor và MaskedTensor . Bộ trang trí công văn có thể được sử dụng để ghi đè hành vi mặc định của các hoạt động TensorFlow khi áp dụng cho các loại phần mở rộng giống tensor. TensorFlow hiện xác định ba trình trang trí điều phối:

Gửi cho một API duy nhất

Trình trang trí tf.experimental.dispatch_for_api ghi đè hành vi mặc định của một hoạt động TensorFlow được chỉ định khi nó được gọi với chữ ký được chỉ định. Ví dụ: bạn có thể sử dụng trình trang trí này để chỉ định cách tf.stack sẽ xử lý các giá trị MaskedTensor :

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

Điều này ghi đè triển khai mặc định cho tf.stack bất cứ khi nào nó được gọi với danh sách các giá trị MaskedTensor (vì đối số values được chú thích bằng 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]]>

Để cho phép tf.stack xử lý danh sách các giá trị MaskedTensor và Tensor hỗn hợp, bạn có thể tinh chỉnh chú thích kiểu cho tham số values và cập nhật nội dung của hàm một cách thích hợp:

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

Để biết danh sách các API có thể bị ghi đè, hãy xem tài liệu API cho tf.experimental.dispatch_for_api .

Gửi cho tất cả các API nguyên tố một lần

Trình trang trí tf.experimental.dispatch_for_unary_elementwise_apis ghi đè hành vi mặc định của tất cả các hoạt động phần tử một ngôi (chẳng hạn như tf.math.cos ) bất cứ khi nào giá trị của đối số đầu tiên (thường được đặt tên là x ) khớp với loại chú thích x_type . Hàm được trang trí phải có hai đối số:

  • api_func : Một hàm nhận một tham số duy nhất và thực hiện thao tác theo phần tử (ví dụ: tf.abs ).
  • x : Đối số đầu tiên của phép toán theo phần tử.

Ví dụ sau cập nhật tất cả các thao tác theo từng phần tử một để xử lý kiểu 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)

Chức năng này bây giờ sẽ được sử dụng bất cứ khi nào thao tác theo phần tử một ngôi được gọi trên 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]>

Gửi cho nhị phân tất cả các API nguyên tố

Tương tự, tf.experimental.dispatch_for_binary_elementwise_apis có thể được sử dụng để cập nhật tất cả các hoạt động phần tử nhị phân để xử lý kiểu 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], [_, _, _]]>

Để biết danh sách các API phần tử bị ghi đè, hãy xem tài liệu API cho tf.experimental.dispatch_for_unary_elementwise_apis và tf.experimental.dispatch_for_binary_elementwise_apis .

Mở rộng có thể mở rộng

Một ExtensionType có thể truy cập được nếu một phiên bản duy nhất có thể được sử dụng để đại diện cho một loạt giá trị. Thông thường, điều này được thực hiện bằng cách thêm thứ nguyên lô vào tất cả các Tensor lồng nhau. Các API TensorFlow sau đây yêu cầu bất kỳ đầu vào loại tiện ích mở rộng nào đều có thể truy cập được:

Theo mặc định, BatchableExtensionType tạo các giá trị theo lô bằng cách phân lô bất kỳ Tensor , CompositeTensor và ExtensionType được lồng vào nhau nào. Nếu điều này không phù hợp với lớp của bạn, thì bạn sẽ cần sử dụng tf.experimental.ExtensionTypeBatchEncoder để ghi đè hành vi mặc định này. Ví dụ: sẽ không thích hợp để tạo một loạt các giá trị tf.SparseTensor bằng cách chỉ cần xếp chồng các values , indices và các trường dense_shape thưa thớt riêng lẻ - trong hầu hết các trường hợp, bạn không thể xếp chồng các tensors này, vì chúng có hình dạng không tương thích ; và ngay cả khi bạn có thể, kết quả sẽ không phải là một SparseTensor hợp lệ.

Ví dụ về BatchableExtensionType: Mạng

Ví dụ: hãy xem xét một lớp Network đơn giản được sử dụng để cân bằng tải, lớp này theo dõi lượng công việc còn lại phải làm ở mỗi nút và lượng băng thông có sẵn để di chuyển công việc giữa các nút:

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

Để làm cho loại này có thể truy cập được, hãy thay đổi loại cơ sở thành BatchableExtensionType và điều chỉnh hình dạng của từng trường để bao gồm các kích thước lô tùy chọn. Ví dụ sau đây cũng thêm một trường shape để theo dõi hình dạng lô. Trường shape này không được yêu cầu bởi tf.data.Dataset hoặc tf.map_fn , nhưng nó được yêu cầu bởi 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.]]]>

Sau đó, bạn có thể sử dụng tf.data.Dataset để lặp lại qua một loạt mạng:

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

Và bạn cũng có thể sử dụng map_fn để áp dụng một hàm cho từng phần tử lô:

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

Các API TensorFlow hỗ trợ ExtensionTypes

@ tf. function

tf. functions là một trình trang trí tính toán trước các đồ thị TensorFlow cho các hàm Python, có thể cải thiện đáng kể hiệu suất của mã TensorFlow của bạn. Các giá trị kiểu mở rộng có thể được sử dụng một cách rõ ràng với các hàm được trang trí @tf.function . Chức năng.

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

Nếu bạn muốn chỉ định rõ ràng input_signature cho tf. tf.function , thì bạn có thể làm như vậy bằng cách sử dụng TypeSpec của loại phần mở rộng.

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

Chức năng bê tông

Các hàm cụ thể đóng gói các đồ thị theo dấu vết riêng lẻ được xây dựng bởi tf.function . function. Các kiểu mở rộng có thể được sử dụng trong suốt với các chức năng cụ thể.

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

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

Kiểm soát hoạt động của dòng chảy

Các loại tiện ích mở rộng được hỗ trợ bởi các hoạt động luồng điều khiển của 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, _]>

Dòng kiểm soát chữ ký

Các loại phần mở rộng cũng được hỗ trợ bởi các câu lệnh luồng điều khiển trong tf. function (sử dụng chữ ký). Trong ví dụ sau, câu lệnh if và câu lệnh for được tự động chuyển đổi thành các phép tf.cond và tf.while_loop , hỗ trợ các kiểu mở rộng.

@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 là API cấp cao của TensorFlow để xây dựng và đào tạo các mô hình học sâu. Các kiểu mở rộng có thể được chuyển làm đầu vào cho mô hình Keras, được chuyển giữa các lớp Keras và được mô hình Keras trả về. Keras hiện đưa ra hai yêu cầu đối với các loại tiện ích mở rộng:

  • Chúng phải có thể truy cập được (xem "Loại mở rộng có thể mở rộng" ở trên).
  • Hình shape phải có một trường hoặc thuộc tính được đặt tên. shape[0] được giả định là thứ nguyên của lô.

Hai phần phụ sau đây đưa ra các ví dụ cho thấy cách sử dụng các loại tiện ích mở rộng với Keras.

Ví dụ về Keras: Network

Đối với ví dụ đầu tiên, hãy xem xét lớp Network được định nghĩa trong phần "Loại mở rộng có thể mở rộng" ở trên, lớp này có thể được sử dụng cho công việc cân bằng tải giữa các nút. Định nghĩa của nó được lặp lại ở đây:

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

Bạn có thể xác định một lớp Keras mới xử lý các Network .

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

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

Sau đó, bạn có thể sử dụng các lớp này để tạo một mô hình đơn giản. Để nạp một ExtensionType vào một mô hình, bạn có thể sử dụng lớp tf.keras.layer.Input với type_spec được đặt thành TypeSpec của loại tiện ích mở rộng. Nếu mô hình Keras sẽ được sử dụng để xử lý lô, thì type_spec phải bao gồm thứ nguyên lô.

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

Cuối cùng, bạn có thể áp dụng mô hình cho một mạng đơn lẻ và cho một loạt mạng.

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

Ví dụ về Keras: MaskedTensor

Trong ví dụ này, MaskedTensor được mở rộng để hỗ trợ Keras . shape được định nghĩa là một thuộc tính được tính toán từ trường values . Keras yêu cầu bạn thêm thuộc tính này vào cả kiểu mở rộng và TypeSpec của nó. MaskedTensor cũng định nghĩa một biến __name__ , biến này sẽ được yêu cầu để tuần tự hóa SavedModel (bên dưới).

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

Tiếp theo, trình trang trí điều phối được sử dụng để ghi đè hành vi mặc định của một số API TensorFlow. Vì các API này được sử dụng bởi các lớp Keras tiêu chuẩn (chẳng hạn như lớp Dense ), việc ghi đè các API này sẽ cho phép chúng tôi sử dụng các lớp đó với MaskedTensor . Đối với mục đích của ví dụ này, matmul cho bộ căng mặt nạ được định nghĩa để coi các giá trị bị che là số không (nghĩa là không đưa chúng vào sản phẩm).

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

Sau đó, bạn có thể xây dựng mô hình Keras chấp nhận đầu vào MaskedTensor , sử dụng các lớp Keras tiêu chuẩn:

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 là một chương trình TensorFlow được tuần tự hóa, bao gồm cả trọng số và tính toán. Nó có thể được xây dựng từ một mô hình Keras hoặc từ một mô hình tùy chỉnh. Trong cả hai trường hợp, các loại tiện ích mở rộng có thể được sử dụng một cách minh bạch với các chức năng và phương pháp được SavedModel xác định.

SavedModel có thể lưu các mô hình, lớp và chức năng xử lý các loại tiện ích mở rộng, miễn là các loại tiện ích mở rộng có trường __name__ . Tên này được sử dụng để đăng ký kiểu mở rộng, vì vậy nó có thể được định vị khi mô hình được tải.

Ví dụ: lưu một mô hình Keras

Các mô hình Keras sử dụng loại tiện ích mở rộng có thể được lưu bằng 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)>

Ví dụ: lưu một mô hình tùy chỉnh

SavedModel cũng có thể được sử dụng để lưu các lớp con tf.Module tùy chỉnh với các hàm xử lý các loại phần mở rộng.

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

Đang tải SavedModel khi ExtensionType không khả dụng

Nếu bạn tải SavedModel sử dụng ExtensionType , nhưng ExtensionType đó không có sẵn (tức là chưa được nhập), thì bạn sẽ thấy cảnh báo và TensorFlow sẽ quay lại sử dụng đối tượng "loại tiện ích mở rộng ẩn danh". Đối tượng này sẽ có các trường giống như kiểu gốc, nhưng sẽ thiếu bất kỳ tùy chỉnh nào khác mà bạn đã thêm cho kiểu, chẳng hạn như các phương thức hoặc thuộc tính tùy chỉnh.

Sử dụng ExtensionTypes với phục vụ TensorFlow

Hiện tại, phục vụ TensorFlow (và những người tiêu dùng khác của từ điển "chữ ký" của SavedModel) yêu cầu tất cả các đầu vào và đầu ra phải là tensor thô. Nếu bạn muốn sử dụng phân phát TensorFlow với một mô hình sử dụng các loại tiện ích mở rộng, thì bạn có thể thêm các phương thức trình bao bọc để soạn hoặc phân tách các giá trị loại tiện ích mở rộng khỏi các tensor. Ví dụ:

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>

Bộ dữ liệu

tf.data là một API cho phép bạn xây dựng các đường ống đầu vào phức tạp từ các mảnh đơn giản, có thể tái sử dụng. Cấu trúc dữ liệu cốt lõi của nó là tf.data.Dataset , đại diện cho một chuỗi các phần tử, trong đó mỗi phần tử bao gồm một hoặc nhiều thành phần.

Xây dựng tập dữ liệu với các loại tiện ích mở rộng

Tập dữ liệu có thể được tạo từ các giá trị loại tiện ích mở rộng bằng cách sử dụng Dataset.from_tensors , Dataset.from_tensor_slices hoặc 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]>

Kết hợp và bỏ kết hợp Tập dữ liệu với các loại tiện ích mở rộng

Tập dữ liệu với các loại tiện ích mở rộng có thể được truy cập và không được so khớp bằng cách sử dụng 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]>