Giúp bảo vệ Great Barrier Reef với TensorFlow trên Kaggle Tham Challenge

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

Xem trên TensorFlow.org Xem nguồn trên GitHubTải xuống sổ ghi chép

Cài đặt

!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ả API cấp cao (như Keras , tf.function , tf.SavedModel ) và hạ thấp cấp API (như tf.while_looptf.concat ). Loại mở rộng TensorFlow thể được sử dụng để tạo ra các loại đối tượng theo định hướng người dùng xác định rằng công việc liền mạch với các API TensorFlow của. Để tạo một kiểu mở rộng, chỉ cần xác định một lớp học với Python tf.experimental.ExtensionType như cơ sở của nó, và sử dụng các chú thích loại để xác định kiểu cho từng lĩnh vực.

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

Các tf.experimental.ExtensionType lớp cơ sở hoạt động tương tự để typing.NamedTuple@dataclasses.dataclass từ thư viện Python chuẩn. Đặc biệt, nó sẽ tự động thêm một constructor và phương pháp đặc biệt (chẳng hạn như __repr____eq__ ) dựa trên các chú thích loại 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, trong đó nhóm lại với nhau một tập hợp các giá trị liên quan, và có thể cung cấp các hoạt động hữu ích dựa trên những giá trị đó. Cấu trúc dữ liệu có thể được khá chung (chẳng hạn như TensorGraph ví dụ trên); hoặc chúng có thể được tùy chỉnh cao cho một mô hình cụ thể.

  • Tensor giống như các loại, chuyên hoặc mở rộng các khái niệm về "tensor." Loại trong thể loại này có một rank , một shape , và thường là một dtype ; và nó làm cho tinh thần để sử dụng chúng với các hoạt động tensor (như tf.stack , tf.add , hoặc tf.matmul ). MaskedTensorCSRSparseMatrix là những ví dụ của các loại tensor-như thế nào.

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: loại mở rộng có thể được sử dụng như nguyên liệu đầu vào và đầu ra cho Keras ModelsLayers .
  • tf.data.Dataset: loại mở rộng có thể được bao gồm trong Datasets , và trả về bởi bộ dữ liệu Iterators .
  • Hub Tensorflow: loại mở rộng có thể được sử dụng như nguyên liệu đầu vào và đầu ra cho tf.hub mô-đun.
  • SavedModel: loại mở rộng có thể được sử dụng như nguyên liệu đầu vào và đầu ra cho SavedModel chức năng.
  • tf.function: loại mở rộng có thể được sử dụng như các đối số và giá trị trả lại cho các chức năng được bao bọc bởi các @tf.function trang trí.
  • vòng lặp while: loại mở rộng có thể được sử dụng như biến vòng lặp trong tf.while_loop , và có thể được sử dụng như các đối số và giá trị trả lại cho cơ thể của vòng lặp while của.
  • điều kiện: loại mở rộng có thể được lựa chọn có điều kiện sử dụng tf.condtf.case .
  • py_function: loại mở rộng có thể được sử dụng như các đối số và giá trị trả lại cho func lập luận để tf.py_function .
  • Tensor ops: loại mở rộng có thể được mở rộng để hỗ trợ hầu hết ops TensorFlow chấp nhận tensor đầu vào (ví dụ, tf.matmul , tf.gather , và tf.reduce_sum ). Hãy xem phần "văn" dưới đây để biết thêm thông tin.
  • chiến lược phân phối: Các loại mở rộng có thể được sử dụng như giá trị 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ợ:

Kiểu 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 một set_mask phương pháp để biến một MaskedTensor , bạn có thể định nghĩa một replace_mask phương pháp mà lợi nhuận mới 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)

Tính năng bổ sung bởi ExtensionType

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

  • Một constructor ( __init__ ).
  • Một phương pháp đại diện có thể in được ( __repr__ ).
  • Bình đẳng và các nhà khai thác sự bất bình đẳng ( __eq__ ).
  • Một phương pháp xác nhận ( __validate__ ).
  • Tính bất biến bắt buộc.
  • Một lồng TypeSpec .
  • 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

Các nhà xây dựng bổ sung bởi ExtensionType mất từng lĩnh vực như một đố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, Tensor lĩnh vực được chuyển đổi sử dụng tf.convert_to_tensor ; Tuple lĩnh vực được chuyển đổi thành tuple s; và Mapping lĩnh vực được chuyển đổi thành dicts không thay đổi.

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)

Các nhà xây dựng đặt ra một TypeError nếu một giá trị trường không thể được chuyển đổi sang kiểu được 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 mặc định phương pháp đại diện có thể in được ( __repr__ ) bao gồm các tên lớp và giá trị cho từng lĩnh vực:

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 nhà khai thác bình đẳng mặc định ( __eq____ne__ ) mà xem xét hai giá trị tương đương nếu chúng có cùng loại và tất cả các lĩnh vực của họ đều bình đẳng. 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 __validate__ phương pháp, có thể được overriden để thực hiện kiểm tra xác nhận trên các lĩnh vực. 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.

ông theo gương cập nhật MaskedTensor để xác nhận các shape s và dtype s các lĩnh vực 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 đè __setattr____delattr__ phương pháp để ngăn ngừa độ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 ExtensionType lớp có một tương ứng TypeSpec lớp, được tạo ra tự động và lưu dưới dạng <extension_type_name>.Spec .

Lớp này bắt tất cả các thông tin từ một giá trị ngoại trừ các giá trị của bất kỳ tensors lồng nhau. Đặc biệt, TypeSpec cho một giá trị được tạo ra bằng cách thay thế bất kỳ lồng nhau tensor, ExtensionType, hoặc CompositeTensor với nó TypeSpec .

class Player(tf.experimental.ExtensionType):
  name: tf.Tensor
  attributes: Mapping[str, tf.Tensor]

anne = Player("Anne", {"height": 8.3, "speed": 28.1})
anne_spec = tf.type_spec_from_value(anne)
print(anne_spec.name)  # Records dtype and shape, but not the string value.
print(anne_spec.attributes)  # Records keys and TensorSpecs for values.
WARNING:tensorflow:Mapping types may not work well with tf.nest. Prefer using MutableMapping for <class 'tensorflow.python.framework.immutable_dict.ImmutableDict'>
TensorSpec(shape=(), dtype=tf.string, name=None)
ImmutableDict({'height': TensorSpec(shape=(), dtype=tf.float32, name=None), 'speed': TensorSpec(shape=(), dtype=tf.float32, name=None)})

TypeSpec giá trị có thể được xây dựng một cách rõ ràng, hoặc chúng có thể được xây dựng từ một ExtensionType giá trị 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 s được sử dụng bởi TensorFlow với các giá trị được chia thành một thành phần tĩnh và một thành phần năng động:

  • Các thành phần tĩnh (được cố định vào thời điểm đồ thị-xây dựng) được mã hóa với một tf.TypeSpec .
  • Các thành phần năng động (có thể thay đổi mỗi lần đồ thị được điều hành) được mã hóa như một danh sách tf.Tensor s.

Ví dụ, tf.function thoái lui chức năng bọc của mình bất cứ khi nào một cuộc tranh cãi có trước đây chưa từng 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>}))

Để biết thêm thông tin, vui lòng xem tf.function Hướng dẫn .

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 đè mặc định đại diện có thể in được ( __repr__ ).
  • Xác định các phương pháp.
  • Xác định classmethods và staticmethods.
  • Xác định thuộc tính.
  • Ghi đè lên các constructor mặc định ( __init__ ).
  • Ghi đè các nhà điều hành bình đẳng mặc định ( __eq__ ).
  • Xác định các nhà khai thác (như __add____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 MaskedTensor lớp để tạo ra một chuỗi đại diện dễ đọc hơn khi giá trị này được in trong chế độ háo hức.

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ụ, MaskedTensor loại có thể định nghĩa một with_default phương pháp mà trả về một bản sao của self với các giá trị đeo mặt nạ thay thế bằng một trao default giá trị. Phương pháp tùy chọn có thể được chú thích với các @tf.function trang trí.

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

Loại phần mở rộng có thể xác định phương pháp sử dụng @classmethod@staticmethod trang trí. Ví dụ, MaskedTensor loại có thể định nghĩa một phương pháp nhà máy mà mặt nạ bất kỳ phần tử với một 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

Loại phần mở rộng có thể xác định tính bằng cách sử dụng @property trang trí, giống như bất kỳ lớp Python bình thường. Ví dụ, MaskedTensor loại có thể định nghĩa một dtype tài sản đó là một cách viết tắt cho dtype các 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>)

Trọng các nhà điều hành bình đẳng mặc định ( __eq__ )

Bạn có thể ghi đè mặc định __eq__ điều hành với nhiều loại phần mở rộng. Ví dụ sau cập nhật MaskedTensor để bỏ qua các yếu tố đeo mặt nạ khi so sánh 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 các children lĩnh vực vì Node loại chưa được (hoàn toàn) quy định được nêu ra.

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 tiện ích mở rộng có thể được phân lớp bằ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 cơ bản TensorGraph lớp đó sử dụng ba Tensor lĩnh vực để 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 có thêm một Tensor lĩnh vực để ghi lại một "giá trị đặc trư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 của ExtensionType TypeSpec

Mỗi ExtensionType lớp có một tương ứng TypeSpec lớp, được tạo ra tự động và lưu 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 các TypeSpec , chỉ cần xác định lớp lồng nhau của riêng bạn tên Spec , và ExtensionType sẽ sử dụng như là cơ sở cho việc xây dựng tự động TypeSpec . Bạn có thể tùy chỉnh Spec lớp theo:

  • 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 các MaskedTensor.Spec lớp để làm cho nó dễ dàng hơn để sử dụng:

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

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

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

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

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

Này ghi đè việc thực hiện mặc định cho tf.stack bất cứ khi nào nó được gọi với một danh sách các MaskedTensor giá trị (kể từ khi values tham số được chú thích với 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 hỗn hợp MaskedTensorTensor giá trị, bạn có thể tinh chỉnh các loại chú thích cho các values tham số và cập nhật các cơ quan chức năng 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, _]]>

Đối với một danh sách các API có thể được ghi đè, 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

Các tf.experimental.dispatch_for_unary_elementwise_apis trang trí ghi đè hành vi mặc định của tất cả các ops elementwise unary (như tf.math.cos ) bất cứ khi nào giá trị cho đối số đầu tiên (thường có tên x ) phù hợ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 chức năng mà phải mất một tham số duy nhất và thực hiện các hoạt động elementwise (ví dụ, tf.abs ).
  • x : Đối số đầu tiên đến hoạt động elementwise.

Ví dụ sau sẽ cập nhật toàn bộ hoạt động elementwise unary để xử lý các MaskedTensor loại:

@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 một hoạt động elementwise unary được gọi là 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ự như vậy, 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 elementwise nhị phân để xử lý các MaskedTensor loại:

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

Đối với một danh sách các API elementwise được ghi đè, xem tài liệu API cho tf.experimental.dispatch_for_unary_elementwise_apistf.experimental.dispatch_for_binary_elementwise_apis .

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

Một ExtensionTypebatchable nếu một trường hợp duy nhất có thể được sử dụng để đại diện cho một loạt các giá trị. Thông thường, điều này được thực hiện bằng cách thêm kích thước hàng loạt cho tất cả các lồng Tensor s. 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 ra giá trị đồng loạt bởi Trạm trộn bất kỳ lồng nhau Tensor s, CompositeTensor s, và ExtensionType s. Nếu đây là không thích hợp cho lớp học của bạn, sau đó bạn sẽ cần phải sử dụng tf.experimental.ExtensionTypeBatchEncoder thay đổi hành vi mặc định này. Ví dụ, nó sẽ không thích hợp để tạo ra một loạt các tf.SparseTensor giá trị bằng cách đơn giản xếp cá nhân thưa thớt tensors' values , indices , và dense_shape lĩnh vực - trong nhiều trường hợp, bạn không thể ngăn xếp những tensors, vì chúng có hình dạng không tương thích ; và thậm chí nếu bạn có thể, kết quả sẽ không phải là một giá trị SparseTensor .

Ví dụ về BatchableExtensionType: Mạng

Như một ví dụ, hãy xem xét một đơn giản Network lớp được sử dụng để cân bằng tải, trong đó ca khúc bao nhiêu công việc còn lại để làm tại mỗi nút, và bao nhiêu 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]])

Để thực hiện kiểu batchable này, thay đổi kiểu căn cứ để BatchableExtensionType , và điều chỉnh hình dạng của từng lĩnh vực bao gồm kích thước hàng loạt tùy chọn. Ví dụ sau đây cũng cho biết thêm một shape lĩnh vực để theo dõi keept của hình hàng loạt. Đây shape lĩnh vực là không cần thiết bởi tf.data.Dataset hoặc tf.map_fn , nhưng nó cần thiết 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 thông qua một loạt các 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 chức năng để mỗi phần tử hàng loạt:

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.function là một trang trí mà precomputes TensorFlow đồ thị cho các chức năng Python, có thể cải thiện đáng kể hiệu suất của mã TensorFlow của bạn. Giá trị kiểu mở rộng có thể được sử dụng một cách minh bạch với @tf.function chức năng -decorated.

class Pastry(tf.experimental.ExtensionType):
  sweetness: tf.Tensor  # 2d embedding that encodes sweetness
  chewiness: tf.Tensor  # 2d embedding that encodes chewiness

@tf.function
def combine_pastry_features(x: Pastry):
  return (x.sweetness + x.chewiness) / 2

cookie = Pastry(sweetness=[1.2, 0.4], chewiness=[0.8, 0.2])
combine_pastry_features(cookie)
<tf.Tensor: shape=(2,), dtype=float32, numpy=array([1. , 0.3], dtype=float32)>

Nếu bạn muốn xác định một cách rõ ràng input_signature cho tf.function , sau đó bạn có thể làm như vậy bằng cách sử dụng kiểu mở rộng của 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)>)

Chức năng bê tông

Chức năng bê tông đóng gói đồ thị truy cá nhân được xây dựng bởi tf.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, if tuyên bố và for báo cáo sẽ được tự động chuyển đổi sang tf.condtf.while_loop hoạt động, trong đó các loại hỗ trợ 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 TensorFlow của để xây dựng và đào tạo mô hình sâu học tập. 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).
  • Phải có một trường hoặc tài sản mang tên shape . shape[0] được giả định là kích thước hàng loạt.

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.

Keras dụ: Network

Đối với ví dụ đầu tiên, hãy xem xét Network lớp được định nghĩa trong phần "Batchable ExtensionTypes" ở trên, có thể được sử dụng để cân bằng tải công việc 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ể định nghĩa một lớp Keras mới xử lý Network s.

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

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

Sau đó, bạn có thể sử dụng các lớp này để tạo một mô hình đơn giản. Để nuôi một ExtensionType vào một mô hình, bạn có thể sử dụng một tf.keras.layer.Input lớp với type_spec thiết lập để loại phần mở rộng của TypeSpec . Nếu mô hình Keras sẽ được sử dụng để theo lô trình, sau đó các type_spec phải bao gồm kích thước hàng loạt.

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 tài sản đó được tính từ values trường. Keras đòi hỏi thatyou thêm khách sạn này cho cả hai loại mở rộng và nó TypeSpec . MaskedTensor cũng định nghĩa một __name__ biến, mà sẽ được yêu cầu cho SavedModel serialization (dưới đây).

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. Kể từ khi các API này được sử dụng bởi lớp Keras tiêu chuẩn (chẳng hạn như Dense layer), trọng này sẽ cho phép chúng ta sử dụng những lớp với MaskedTensor . Theo mục đích của ví dụ này, matmul cho tensors đeo mặt nạ được xác định để điều trị các giá trị chèn thêm các số không (ví dụ, để không bao gồm chúng trong 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ột mô hình Keras chấp nhận MaskedTensor đầu vào, sử dụng 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

Một SavedModel là một chương trình TensorFlow serialized, bao gồm cả trọng lượng 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ể tiết kiệm các mô hình, các lớp, và chức năng mà các loại quá trình gia hạn, miễn là các loại phần mở rộng có một __name__ lĩnh vực. 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

Mô hình Keras rằng sử dụng các loại phần 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 tùy chỉnh tf.Module lớp con với chức năng rằng các loại quá trình 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 một SavedModel có sử dụng một ExtensionType , nhưng điều đó ExtensionType không có sẵn (ví dụ, chưa được nhập khẩu), sau đó bạn sẽ thấy một cảnh báo và TensorFlow sẽ quay lại sử dụng một "kiểu mở rộng vô danh" đối tượng. Đố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 nay, TensorFlow phục vụ (và người tiêu dùng khác của SavedModel "chữ ký" từ điển) yêu cầu tất cả đầu vào và đầu ra được tensors 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ừ đơn giản, miếng 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 yếu 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

Datasets thể được xây dựng từ các giá trị kiểu mở rộng 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à hủy ghép các Tập dữ liệu với các loại tiện ích mở rộng

Datasets với các loại phần mở rộng có thể batchand và unbatched 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]>