Halaman ini diterjemahkan oleh Cloud Translation API.

Bermigrasi dari TPU embedding_columns ke lapisan TPUEbedding

Lihat di TensorFlow.org

Jalankan di Google Colab

Lihat sumber di GitHub

Unduh buku catatan

Panduan ini menunjukkan cara memigrasikan pelatihan penyematan pada TPU dari API embedding_column TensorFlow 1 dengan TPUEstimator ke API lapisan TPUEmbedding TensorFlow 2 dengan TPUStrategy .

Embeddings adalah matriks (besar). Mereka adalah tabel pencarian yang memetakan dari ruang fitur yang jarang ke vektor padat. Embeddings memberikan representasi yang efisien dan padat, menangkap kesamaan kompleks dan hubungan antar fitur.

TensorFlow menyertakan dukungan khusus untuk pelatihan penyematan pada TPU. Dukungan penyematan khusus TPU ini memungkinkan Anda untuk melatih penyematan yang lebih besar dari memori satu perangkat TPU, dan menggunakan input yang jarang dan tidak rata pada TPU.

Di TensorFlow 1, tf.compat.v1.estimator.tpu.TPUEstimator adalah API tingkat tinggi yang merangkum pelatihan, evaluasi, prediksi, dan ekspor untuk disajikan dengan TPU. Ini memiliki dukungan khusus untuk tf.compat.v1.tpu.experimental.embedding_column .
Untuk menerapkan ini di TensorFlow 2, gunakan lapisan tfrs.layers.embedding.TPUEmbedding dari TensorFlow Recommenders. Untuk pelatihan dan evaluasi, gunakan strategi distribusi tf.distribute.TPUStrategy —yang kompatibel dengan Keras API untuk, misalnya, pembuatan model ( tf.keras.Model ), pengoptimal ( tf.keras.optimizers.Optimizer ), dan pelatihan dengan Model.fit atau loop pelatihan khusus dengan tf.function dan tf.GradientTape .

Untuk informasi tambahan, lihat dokumentasi API lapisan tfrs.layers.embedding.TPUEmbedding , serta dokumen tf.tpu.experimental.embedding.TableConfig dan tf.tpu.experimental.embedding.FeatureConfig untuk informasi tambahan. Untuk gambaran umum tentang tf.distribute.TPUStrategy , lihat panduan pelatihan Terdistribusi dan panduan Penggunaan TPU . Jika Anda bermigrasi dari TPUEstimator ke TPUStrategy , lihat panduan migrasi TPU .

Mempersiapkan

Mulailah dengan menginstal TensorFlow Recommenders dan mengimpor beberapa paket yang diperlukan:

pip install tensorflow-recommenders

import tensorflow as tf
import tensorflow.compat.v1 as tf1

# TPUEmbedding layer is not part of TensorFlow.
import tensorflow_recommenders as tfrs

/tmpfs/src/tf_docs_env/lib/python3.6/site-packages/requests/__init__.py:104: RequestsDependencyWarning: urllib3 (1.26.8) or chardet (2.3.0)/charset_normalizer (2.0.11) doesn't match a supported version!
  RequestsDependencyWarning)

Dan siapkan dataset sederhana untuk tujuan demonstrasi:

features = [[1., 1.5]]
embedding_features_indices = [[0, 0], [0, 1]]
embedding_features_values = [0, 5]
labels = [[0.3]]
eval_features = [[4., 4.5]]
eval_embedding_features_indices = [[0, 0], [0, 1]]
eval_embedding_features_values = [4, 3]
eval_labels = [[0.8]]

TensorFlow 1: Latih penyematan pada TPU dengan TPUEstimator

Di TensorFlow 1, Anda menyiapkan penyematan TPU menggunakan tf.compat.v1.tpu.experimental.embedding_column API dan melatih/mengevaluasi model pada TPU dengan tf.compat.v1.estimator.tpu.TPUEstimator .

Inputnya adalah bilangan bulat mulai dari nol hingga ukuran kosakata untuk tabel penyematan TPU. Mulailah dengan menyandikan input ke ID kategoris dengan tf.feature_column.categorical_column_with_identity . Gunakan "sparse_feature" untuk parameter key , karena fitur input bernilai integer, sedangkan num_buckets adalah ukuran kosakata untuk tabel embedding ( 10 ).

embedding_id_column = (
      tf1.feature_column.categorical_column_with_identity(
          key="sparse_feature", num_buckets=10))

Selanjutnya, ubah input kategorikal sparse menjadi representasi padat dengan tpu.experimental.embedding_column , di mana dimension adalah lebar dari tabel embedding. Ini akan menyimpan vektor embedding untuk masing-masing num_buckets .

embedding_column = tf1.tpu.experimental.embedding_column(
    embedding_id_column, dimension=5)

Sekarang, tentukan konfigurasi penyematan khusus TPU melalui tf.estimator.tpu.experimental.EmbeddingConfigSpec . Anda akan meneruskannya nanti ke tf.estimator.tpu.TPUEstimator sebagai parameter embedding_config_spec .

embedding_config_spec = tf1.estimator.tpu.experimental.EmbeddingConfigSpec(
    feature_columns=(embedding_column,),
    optimization_parameters=(
        tf1.tpu.experimental.AdagradParameters(0.05)))

Selanjutnya, untuk menggunakan TPUEstimator , tentukan:

Fungsi input untuk data pelatihan
Fungsi input evaluasi untuk data evaluasi
Fungsi model untuk menginstruksikan TPUEstimator bagaimana operasi pelatihan didefinisikan dengan fitur dan label

def _input_fn(params):
  dataset = tf1.data.Dataset.from_tensor_slices((
      {"dense_feature": features,
       "sparse_feature": tf1.SparseTensor(
           embedding_features_indices,
           embedding_features_values, [1, 2])},
           labels))
  dataset = dataset.repeat()
  return dataset.batch(params['batch_size'], drop_remainder=True)

def _eval_input_fn(params):
  dataset = tf1.data.Dataset.from_tensor_slices((
      {"dense_feature": eval_features,
       "sparse_feature": tf1.SparseTensor(
           eval_embedding_features_indices,
           eval_embedding_features_values, [1, 2])},
           eval_labels))
  dataset = dataset.repeat()
  return dataset.batch(params['batch_size'], drop_remainder=True)

def _model_fn(features, labels, mode, params):
  embedding_features = tf1.keras.layers.DenseFeatures(embedding_column)(features)
  concatenated_features = tf1.keras.layers.Concatenate(axis=1)(
      [embedding_features, features["dense_feature"]])
  logits = tf1.layers.Dense(1)(concatenated_features)
  loss = tf1.losses.mean_squared_error(labels=labels, predictions=logits)
  optimizer = tf1.train.AdagradOptimizer(0.05)
  optimizer = tf1.tpu.CrossShardOptimizer(optimizer)
  train_op = optimizer.minimize(loss, global_step=tf1.train.get_global_step())
  return tf1.estimator.tpu.TPUEstimatorSpec(mode, loss=loss, train_op=train_op)

Dengan fungsi tersebut didefinisikan, buat tf.distribute.cluster_resolver.TPUClusterResolver yang menyediakan informasi cluster, dan objek tf.compat.v1.estimator.tpu.RunConfig .

Seiring dengan fungsi model yang telah Anda tetapkan, Anda sekarang dapat membuat TPUEstimator . Di sini, Anda akan menyederhanakan alur dengan melewatkan penghematan pos pemeriksaan. Kemudian, Anda akan menentukan ukuran batch untuk pelatihan dan evaluasi untuk TPUEstimator .

cluster_resolver = tf1.distribute.cluster_resolver.TPUClusterResolver(tpu='')
print("All devices: ", tf1.config.list_logical_devices('TPU'))

All devices:  []

tpu_config = tf1.estimator.tpu.TPUConfig(
    iterations_per_loop=10,
    per_host_input_for_training=tf1.estimator.tpu.InputPipelineConfig
          .PER_HOST_V2)
config = tf1.estimator.tpu.RunConfig(
    cluster=cluster_resolver,
    save_checkpoints_steps=None,
    tpu_config=tpu_config)
estimator = tf1.estimator.tpu.TPUEstimator(
    model_fn=_model_fn, config=config, train_batch_size=8, eval_batch_size=8,
    embedding_config_spec=embedding_config_spec)

WARNING:tensorflow:Estimator's model_fn (<function _model_fn at 0x7eff1dbf4ae8>) includes params argument, but params are not passed to Estimator.
WARNING:tensorflow:Using temporary folder as model directory: /tmp/tmpc68an8jx
INFO:tensorflow:Using config: {'_model_dir': '/tmp/tmpc68an8jx', '_tf_random_seed': None, '_save_summary_steps': 100, '_save_checkpoints_steps': None, '_save_checkpoints_secs': None, '_session_config': allow_soft_placement: true
cluster_def {
  job {
    name: "worker"
    tasks {
      key: 0
      value: "10.240.1.2:8470"
    }
  }
}
isolate_session_state: true
, '_keep_checkpoint_max': 5, '_keep_checkpoint_every_n_hours': 10000, '_log_step_count_steps': None, '_train_distribute': None, '_device_fn': None, '_protocol': None, '_eval_distribute': None, '_experimental_distribute': None, '_experimental_max_worker_delay_secs': None, '_session_creation_timeout_secs': 7200, '_checkpoint_save_graph_def': True, '_service': None, '_cluster_spec': ClusterSpec({'worker': ['10.240.1.2:8470']}), '_task_type': 'worker', '_task_id': 0, '_global_id_in_cluster': 0, '_master': 'grpc://10.240.1.2:8470', '_evaluation_master': 'grpc://10.240.1.2:8470', '_is_chief': True, '_num_ps_replicas': 0, '_num_worker_replicas': 1, '_tpu_config': TPUConfig(iterations_per_loop=10, num_shards=None, num_cores_per_replica=None, per_host_input_for_training=3, tpu_job_name=None, initial_infeed_sleep_secs=None, input_partition_dims=None, eval_training_input_configuration=2, experimental_host_call_every_n_steps=1, experimental_allow_per_host_v2_parallel_get_next=False, experimental_feed_hook=None), '_cluster': <tensorflow.python.distribute.cluster_resolver.tpu.tpu_cluster_resolver.TPUClusterResolver object at 0x7eff1dbfa2b0>}
INFO:tensorflow:_TPUContext: eval_on_tpu True

Hubungi TPUEstimator.train untuk mulai melatih model:

estimator.train(_input_fn, steps=1)

INFO:tensorflow:Querying Tensorflow master (grpc://10.240.1.2:8470) for TPU system metadata.
INFO:tensorflow:Found TPU system:
INFO:tensorflow:*** Num TPU Cores: 8
INFO:tensorflow:*** Num TPU Workers: 1
INFO:tensorflow:*** Num TPU Cores Per Worker: 8
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:CPU:0, CPU, -1, -3018931587863375246)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:0, TPU, 17179869184, 1249032734884062775)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:1, TPU, 17179869184, -3881759543008185868)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:2, TPU, 17179869184, -3421771184935649663)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:3, TPU, 17179869184, 8872583169621331661)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:4, TPU, 17179869184, -1222373804129613329)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:5, TPU, 17179869184, 6258068298163390748)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:6, TPU, 17179869184, 5190265587768274342)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:7, TPU, 17179869184, 3073578684150069836)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU_SYSTEM:0, TPU_SYSTEM, 17179869184, 2071242092327503173)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 17179869184, -1319360343564144287)
WARNING:tensorflow:From /tmpfs/src/tf_docs_env/lib/python3.6/site-packages/tensorflow/python/training/training_util.py:236: Variable.initialized_value (from tensorflow.python.ops.variables) is deprecated and will be removed in a future version.
Instructions for updating:
Use Variable.read_value. Variables in 2.X are initialized automatically both in eager and graph (inside tf.defun) contexts.
INFO:tensorflow:Calling model_fn.
WARNING:tensorflow:From /tmpfs/src/tf_docs_env/lib/python3.6/site-packages/tensorflow/python/tpu/feature_column_v2.py:479: IdentityCategoricalColumn._num_buckets (from tensorflow.python.feature_column.feature_column_v2) is deprecated and will be removed in a future version.
Instructions for updating:
The old _FeatureColumn APIs are being deprecated. Please use the new FeatureColumn APIs instead.
INFO:tensorflow:Querying Tensorflow master (grpc://10.240.1.2:8470) for TPU system metadata.
INFO:tensorflow:Found TPU system:
INFO:tensorflow:*** Num TPU Cores: 8
INFO:tensorflow:*** Num TPU Workers: 1
INFO:tensorflow:*** Num TPU Cores Per Worker: 8
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:CPU:0, CPU, -1, -3018931587863375246)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:0, TPU, 17179869184, 1249032734884062775)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:1, TPU, 17179869184, -3881759543008185868)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:2, TPU, 17179869184, -3421771184935649663)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:3, TPU, 17179869184, 8872583169621331661)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:4, TPU, 17179869184, -1222373804129613329)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:5, TPU, 17179869184, 6258068298163390748)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:6, TPU, 17179869184, 5190265587768274342)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:7, TPU, 17179869184, 3073578684150069836)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU_SYSTEM:0, TPU_SYSTEM, 17179869184, 2071242092327503173)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 17179869184, -1319360343564144287)
WARNING:tensorflow:From /tmpfs/src/tf_docs_env/lib/python3.6/site-packages/tensorflow/python/training/adagrad.py:77: calling Constant.__init__ (from tensorflow.python.ops.init_ops) with dtype is deprecated and will be removed in a future version.
Instructions for updating:
Call initializer instance with the dtype argument instead of passing it to the constructor
INFO:tensorflow:Bypassing TPUEstimator hook
INFO:tensorflow:Done calling model_fn.
INFO:tensorflow:TPU job name worker
INFO:tensorflow:Graph was finalized.
INFO:tensorflow:Running local_init_op.
INFO:tensorflow:Done running local_init_op.
WARNING:tensorflow:From /tmpfs/src/tf_docs_env/lib/python3.6/site-packages/tensorflow_estimator/python/estimator/tpu/tpu_estimator.py:758: Variable.load (from tensorflow.python.ops.variables) is deprecated and will be removed in a future version.
Instructions for updating:
Prefer Variable.assign which has equivalent behavior in 2.X.
INFO:tensorflow:Initialized dataset iterators in 0 seconds
INFO:tensorflow:Installing graceful shutdown hook.
INFO:tensorflow:Creating heartbeat manager for ['/job:worker/replica:0/task:0/device:CPU:0']
INFO:tensorflow:Configuring worker heartbeat: shutdown_mode: WAIT_FOR_COORDINATOR

INFO:tensorflow:Init TPU system
INFO:tensorflow:Initialized TPU in 9 seconds
INFO:tensorflow:Starting infeed thread controller.
INFO:tensorflow:Starting outfeed thread controller.
INFO:tensorflow:Enqueue next (1) batch(es) of data to infeed.
INFO:tensorflow:Dequeue next (1) batch(es) of data from outfeed.
INFO:tensorflow:Outfeed finished for iteration (0, 0)
INFO:tensorflow:loss = 0.5212165, step = 1
INFO:tensorflow:Stop infeed thread controller
INFO:tensorflow:Shutting down InfeedController thread.
INFO:tensorflow:InfeedController received shutdown signal, stopping.
INFO:tensorflow:Infeed thread finished, shutting down.
INFO:tensorflow:infeed marked as finished
INFO:tensorflow:Stop output thread controller
INFO:tensorflow:Shutting down OutfeedController thread.
INFO:tensorflow:OutfeedController received shutdown signal, stopping.
INFO:tensorflow:Outfeed thread finished, shutting down.
INFO:tensorflow:outfeed marked as finished
INFO:tensorflow:Shutdown TPU system.
INFO:tensorflow:Loss for final step: 0.5212165.
INFO:tensorflow:training_loop marked as finished
<tensorflow_estimator.python.estimator.tpu.tpu_estimator.TPUEstimator at 0x7eff1dbfa7b8>

Kemudian, panggil TPUEstimator.evaluate untuk mengevaluasi model menggunakan data evaluasi:

estimator.evaluate(_eval_input_fn, steps=1)

INFO:tensorflow:Could not find trained model in model_dir: /tmp/tmpc68an8jx, running initialization to evaluate.
INFO:tensorflow:Calling model_fn.
INFO:tensorflow:Querying Tensorflow master (grpc://10.240.1.2:8470) for TPU system metadata.
INFO:tensorflow:Found TPU system:
INFO:tensorflow:*** Num TPU Cores: 8
INFO:tensorflow:*** Num TPU Workers: 1
INFO:tensorflow:*** Num TPU Cores Per Worker: 8
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:CPU:0, CPU, -1, -3018931587863375246)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:0, TPU, 17179869184, 1249032734884062775)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:1, TPU, 17179869184, -3881759543008185868)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:2, TPU, 17179869184, -3421771184935649663)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:3, TPU, 17179869184, 8872583169621331661)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:4, TPU, 17179869184, -1222373804129613329)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:5, TPU, 17179869184, 6258068298163390748)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:6, TPU, 17179869184, 5190265587768274342)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:7, TPU, 17179869184, 3073578684150069836)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU_SYSTEM:0, TPU_SYSTEM, 17179869184, 2071242092327503173)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 17179869184, -1319360343564144287)
WARNING:tensorflow:From /tmpfs/src/tf_docs_env/lib/python3.6/site-packages/tensorflow_estimator/python/estimator/tpu/tpu_estimator.py:3406: div (from tensorflow.python.ops.math_ops) is deprecated and will be removed in a future version.
Instructions for updating:
Deprecated in favor of operator or tf.math.divide.
INFO:tensorflow:Done calling model_fn.
INFO:tensorflow:Starting evaluation at 2022-02-05T13:21:42
INFO:tensorflow:TPU job name worker
INFO:tensorflow:Graph was finalized.
INFO:tensorflow:Running local_init_op.
INFO:tensorflow:Done running local_init_op.
INFO:tensorflow:Init TPU system
INFO:tensorflow:Initialized TPU in 11 seconds
INFO:tensorflow:Starting infeed thread controller.
INFO:tensorflow:Starting outfeed thread controller.
INFO:tensorflow:Initialized dataset iterators in 0 seconds
INFO:tensorflow:Enqueue next (1) batch(es) of data to infeed.
INFO:tensorflow:Dequeue next (1) batch(es) of data from outfeed.
INFO:tensorflow:Outfeed finished for iteration (0, 0)
INFO:tensorflow:Evaluation [1/1]
INFO:tensorflow:Stop infeed thread controller
INFO:tensorflow:Shutting down InfeedController thread.
INFO:tensorflow:InfeedController received shutdown signal, stopping.
INFO:tensorflow:Infeed thread finished, shutting down.
INFO:tensorflow:infeed marked as finished
INFO:tensorflow:Stop output thread controller
INFO:tensorflow:Shutting down OutfeedController thread.
INFO:tensorflow:OutfeedController received shutdown signal, stopping.
INFO:tensorflow:Outfeed thread finished, shutting down.
INFO:tensorflow:outfeed marked as finished
INFO:tensorflow:Shutdown TPU system.
INFO:tensorflow:Inference Time : 12.50468s
INFO:tensorflow:Finished evaluation at 2022-02-05-13:21:54
INFO:tensorflow:Saving dict for global step 1: global_step = 1, loss = 36.28813
INFO:tensorflow:evaluation_loop marked as finished
{'loss': 36.28813, 'global_step': 1}

TensorFlow 2: Latih penyematan pada TPU dengan Strategi TPUS

Di TensorFlow 2, untuk melatih pekerja TPU, gunakan tf.distribute.TPUStrategy bersama dengan Keras API untuk definisi model dan pelatihan/evaluasi. (Lihat panduan Menggunakan TPU untuk contoh pelatihan lainnya dengan Keras Model.fit dan loop pelatihan khusus (dengan tf.function dan tf.GradientTape ).)

Karena Anda perlu melakukan beberapa pekerjaan inisialisasi untuk terhubung ke kluster jarak jauh dan menginisialisasi pekerja TPU, mulailah dengan membuat TPUClusterResolver untuk memberikan informasi kluster dan terhubung ke kluster. (Pelajari lebih lanjut di bagian inisialisasi TPU dari panduan Menggunakan TPU .)

cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')
tf.config.experimental_connect_to_cluster(cluster_resolver)
tf.tpu.experimental.initialize_tpu_system(cluster_resolver)
print("All devices: ", tf.config.list_logical_devices('TPU'))

INFO:tensorflow:Clearing out eager caches
INFO:tensorflow:Clearing out eager caches
INFO:tensorflow:Initializing the TPU system: grpc://10.240.1.2:8470
INFO:tensorflow:Initializing the TPU system: grpc://10.240.1.2:8470
INFO:tensorflow:Finished initializing TPU system.
INFO:tensorflow:Finished initializing TPU system.
All devices:  [LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:0', device_type='TPU'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:1', device_type='TPU'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:2', device_type='TPU'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:3', device_type='TPU'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:4', device_type='TPU'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:5', device_type='TPU'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:6', device_type='TPU'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:7', device_type='TPU')]

Selanjutnya, siapkan data Anda. Ini mirip dengan cara Anda membuat kumpulan data dalam contoh TensorFlow 1, kecuali fungsi kumpulan data sekarang melewati objek tf.distribute.InputContext daripada dict params . Anda dapat menggunakan objek ini untuk menentukan ukuran batch lokal (dan untuk host mana pipeline ini, sehingga Anda dapat mempartisi data dengan benar).

Saat menggunakan tfrs.layers.embedding.TPUEmbedding API, penting untuk menyertakan opsi drop_remainder=True saat mengelompokkan kumpulan data dengan Dataset.batch , karena TPUEmbedding memerlukan ukuran kumpulan yang tetap.
Selain itu, ukuran batch yang sama harus digunakan untuk evaluasi dan pelatihan jika dilakukan pada perangkat yang sama.
Terakhir, Anda harus menggunakan tf.keras.utils.experimental.DatasetCreator bersama dengan opsi input khusus— experimental_fetch_to_device=False — tf.distribute.InputOptions (yang menyimpan konfigurasi khusus strategi). Ini ditunjukkan di bawah ini:

global_batch_size = 8

def _input_dataset(context: tf.distribute.InputContext):
  dataset = tf.data.Dataset.from_tensor_slices((
      {"dense_feature": features,
       "sparse_feature": tf.SparseTensor(
           embedding_features_indices,
           embedding_features_values, [1, 2])},
           labels))
  dataset = dataset.shuffle(10).repeat()
  dataset = dataset.batch(
      context.get_per_replica_batch_size(global_batch_size),
      drop_remainder=True)
  return dataset.prefetch(2)

def _eval_dataset(context: tf.distribute.InputContext):
  dataset = tf.data.Dataset.from_tensor_slices((
      {"dense_feature": eval_features,
       "sparse_feature": tf.SparseTensor(
           eval_embedding_features_indices,
           eval_embedding_features_values, [1, 2])},
           eval_labels))
  dataset = dataset.repeat()
  dataset = dataset.batch(
      context.get_per_replica_batch_size(global_batch_size),
      drop_remainder=True)
  return dataset.prefetch(2)

input_options = tf.distribute.InputOptions(
    experimental_fetch_to_device=False)

input_dataset = tf.keras.utils.experimental.DatasetCreator(
    _input_dataset, input_options=input_options)

eval_dataset = tf.keras.utils.experimental.DatasetCreator(
    _eval_dataset, input_options=input_options)

Selanjutnya, setelah data disiapkan, Anda akan membuat TPUStrategy , dan menentukan model, metrik, dan pengoptimal di bawah cakupan strategi ini ( Strategy.scope ).

Anda harus memilih nomor untuk steps_per_execution di Model.compile karena ini menentukan jumlah batch yang akan dijalankan selama setiap panggilan tf.function , dan sangat penting untuk kinerja. Argumen ini mirip dengan iterations_per_loop yang digunakan dalam TPUEstimator .

Fitur dan konfigurasi tabel yang ditentukan di TensorFlow 1 melalui tf.tpu.experimental.embedding_column (dan tf.tpu.experimental.shared_embedding_column ) dapat ditentukan langsung di TensorFlow 2 melalui sepasang objek konfigurasi:

(Lihat dokumentasi API terkait untuk detail selengkapnya.)

strategy = tf.distribute.TPUStrategy(cluster_resolver)
with strategy.scope():
  optimizer = tf.keras.optimizers.Adagrad(learning_rate=0.05)
  dense_input = tf.keras.Input(shape=(2,), dtype=tf.float32, batch_size=global_batch_size)
  sparse_input = tf.keras.Input(shape=(), dtype=tf.int32, batch_size=global_batch_size)
  embedded_input = tfrs.layers.embedding.TPUEmbedding(
      feature_config=tf.tpu.experimental.embedding.FeatureConfig(
          table=tf.tpu.experimental.embedding.TableConfig(
              vocabulary_size=10,
              dim=5,
              initializer=tf.initializers.TruncatedNormal(mean=0.0, stddev=1)),
          name="sparse_input"),
      optimizer=optimizer)(sparse_input)
  input = tf.keras.layers.Concatenate(axis=1)([dense_input, embedded_input])
  result = tf.keras.layers.Dense(1)(input)
  model = tf.keras.Model(inputs={"dense_feature": dense_input, "sparse_feature": sparse_input}, outputs=result)
  model.compile(optimizer, "mse", steps_per_execution=10)

INFO:tensorflow:Found TPU system:
INFO:tensorflow:Found TPU system:
INFO:tensorflow:*** Num TPU Cores: 8
INFO:tensorflow:*** Num TPU Cores: 8
INFO:tensorflow:*** Num TPU Workers: 1
INFO:tensorflow:*** Num TPU Workers: 1
INFO:tensorflow:*** Num TPU Cores Per Worker: 8
INFO:tensorflow:*** Num TPU Cores Per Worker: 8
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:localhost/replica:0/task:0/device:CPU:0, CPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:localhost/replica:0/task:0/device:CPU:0, CPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:CPU:0, CPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:CPU:0, CPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:0, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:0, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:1, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:1, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:2, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:2, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:3, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:3, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:4, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:4, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:5, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:5, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:6, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:6, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:7, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:7, TPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU_SYSTEM:0, TPU_SYSTEM, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU_SYSTEM:0, TPU_SYSTEM, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 0, 0)
INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 0, 0)

Dengan itu, Anda siap melatih model dengan set data pelatihan:

model.fit(input_dataset, epochs=5, steps_per_epoch=10)

Epoch 1/5
10/10 [==============================] - 2s 164ms/step - loss: 0.4005
Epoch 2/5
10/10 [==============================] - 0s 3ms/step - loss: 0.0036
Epoch 3/5
10/10 [==============================] - 0s 3ms/step - loss: 3.0932e-05
Epoch 4/5
10/10 [==============================] - 0s 3ms/step - loss: 2.5767e-07
Epoch 5/5
10/10 [==============================] - 0s 3ms/step - loss: 2.1366e-09
<keras.callbacks.History at 0x7efd8c461c18>

Terakhir, evaluasi model menggunakan dataset evaluasi:

model.evaluate(eval_dataset, steps=1, return_dict=True)

1/1 [==============================] - 1s 1s/step - loss: 15.3952
{'loss': 15.395216941833496}

Langkah selanjutnya

Pelajari lebih lanjut tentang menyiapkan penyematan khusus TPU di dokumen API:

tfrs.layers.embedding.TPUEmbedding : khususnya tentang fitur dan konfigurasi tabel, menyetel pengoptimal, membuat model (menggunakan API fungsional Keras atau melalui subkelas tf.keras.Model ), pelatihan/evaluasi, dan penyajian model dengan tf.saved_model
tf.tpu.experimental.embedding.TableConfig
tf.tpu.experimental.embedding.FeatureConfig

Untuk informasi selengkapnya tentang TPUStrategy di TensorFlow 2, pertimbangkan referensi berikut:

Panduan: Gunakan TPU (meliputi pelatihan dengan Keras Model.fit /loop pelatihan khusus dengan tf.distribute.TPUStrategy , serta tip untuk meningkatkan kinerja dengan tf.function )
Panduan: Pelatihan terdistribusi dengan TensorFlow
Panduan: Bermigrasi dari TPUEstimator ke TPUStrategy .

Untuk mempelajari lebih lanjut tentang menyesuaikan pelatihan Anda, lihat:

Panduan: Sesuaikan apa yang terjadi di Model.fit
Panduan: Menulis loop pelatihan dari awal

TPU—ASIC khusus Google untuk machine learning—tersedia melalui Google Colab , TPU Research Cloud , dan Cloud TPU .