Pelatihan multi-pekerja dengan Estimator

Peringatan: Penaksir tidak disarankan untuk kode baru. Penaksir menjalankan kode gaya v1.Session yang lebih sulit untuk ditulis dengan benar, dan dapat berperilaku tidak terduga, terutama bila dikombinasikan dengan kode TF 2. Penaksir memang termasuk dalam jaminan kompatibilitas , tetapi tidak akan menerima perbaikan selain kerentanan keamanan. Lihat panduan migrasi untuk detailnya.

Ringkasan

Catatan: Meskipun Anda dapat menggunakan Estimator dengan tf.distribute API, disarankan untuk menggunakan Keras dengan tf.distribute , lihat pelatihan multi-pekerja dengan Keras . Pelatihan penaksir dengan tf.distribute.Strategy memiliki dukungan terbatas.

Tutorial ini menunjukkan bagaimana tf.distribute.Strategy dapat digunakan untuk pelatihan multi-pekerja terdistribusi dengan tf.estimator . Jika Anda menulis kode menggunakan tf.estimator , dan Anda tertarik untuk menskalakan lebih dari satu mesin dengan kinerja tinggi, tutorial ini cocok untuk Anda.

Sebelum memulai, harap baca panduan strategi distribusi . Tutorial pelatihan multi-GPU juga relevan, karena tutorial ini menggunakan model yang sama.

Mempersiapkan

Pertama, siapkan TensorFlow dan impor yang diperlukan.

import tensorflow_datasets as tfds
import tensorflow as tf

import os, json

Catatan: Mulai dari TF2.4, strategi mirror multi-pekerja gagal dengan estimator jika dijalankan dengan bersemangat diaktifkan (default). Kesalahan dalam TF2.4 adalah TypeError: cannot pickle '_thread.lock' object , Lihat edisi #46556 untuk detailnya. Solusinya adalah menonaktifkan eksekusi yang bersemangat.

tf.compat.v1.disable_eager_execution()

Fungsi masukan

Tutorial ini menggunakan set data MNIST dari TensorFlow Datasets . Kode di sini mirip dengan tutorial pelatihan multi-GPU dengan satu perbedaan utama: saat menggunakan Estimator untuk pelatihan multi-pekerja, kumpulan data perlu di-shard dengan jumlah pekerja untuk memastikan konvergensi model. Data input di-sharding oleh indeks pekerja, sehingga setiap pekerja memproses 1/num_workers bagian berbeda dari kumpulan data.

BUFFER_SIZE = 10000
BATCH_SIZE = 64

def input_fn(mode, input_context=None):
  datasets, info = tfds.load(name='mnist',
                                with_info=True,
                                as_supervised=True)
  mnist_dataset = (datasets['train'] if mode == tf.estimator.ModeKeys.TRAIN else
                   datasets['test'])

  def scale(image, label):
    image = tf.cast(image, tf.float32)
    image /= 255
    return image, label

  if input_context:
    mnist_dataset = mnist_dataset.shard(input_context.num_input_pipelines,
                                        input_context.input_pipeline_id)
  return mnist_dataset.map(scale).cache().shuffle(BUFFER_SIZE).batch(BATCH_SIZE)

Pendekatan lain yang masuk akal untuk mencapai konvergensi adalah dengan mengacak dataset dengan benih yang berbeda di setiap pekerja.

Konfigurasi multi-pekerja

Salah satu perbedaan utama dalam tutorial ini (dibandingkan dengan tutorial pelatihan multi-GPU ) adalah pengaturan multi-pekerja. Variabel lingkungan TF_CONFIG adalah cara standar untuk menentukan konfigurasi cluster untuk setiap pekerja yang merupakan bagian dari cluster.

Ada dua komponen TF_CONFIG : cluster dan task . cluster memberikan informasi tentang seluruh cluster, yaitu pekerja dan server parameter di cluster. task memberikan informasi tentang tugas saat ini. cluster komponen pertama sama untuk semua pekerja dan server parameter di cluster, dan task komponen kedua berbeda pada setiap pekerja dan server parameter dan menentukan type dan index sendiri. Dalam contoh ini, type tugas adalah worker dan index tugas adalah 0 .

Sebagai ilustrasi, tutorial ini menunjukkan cara menyetel TF_CONFIG dengan 2 pekerja di localhost . Dalam praktiknya, Anda akan membuat beberapa pekerja pada alamat IP dan port eksternal, dan menyetel TF_CONFIG pada setiap pekerja dengan tepat, yaitu memodifikasi index tugas .

Peringatan: Jangan jalankan kode berikut di Colab. Waktu proses TensorFlow akan mencoba membuat server gRPC pada alamat IP dan port yang ditentukan, yang kemungkinan besar akan gagal. Lihat versi keras dari tutorial ini untuk contoh bagaimana Anda dapat menguji menjalankan banyak pekerja pada satu mesin.

os.environ['TF_CONFIG'] = json.dumps({
    'cluster': {
        'worker': ["localhost:12345", "localhost:23456"]
    },
    'task': {'type': 'worker', 'index': 0}
})

Tentukan modelnya

Tulis lapisan, pengoptimal, dan fungsi kerugian untuk pelatihan. Tutorial ini mendefinisikan model dengan lapisan Keras, mirip dengan tutorial pelatihan multi-GPU .

LEARNING_RATE = 1e-4
def model_fn(features, labels, mode):
  model = tf.keras.Sequential([
      tf.keras.layers.Conv2D(32, 3, activation='relu', input_shape=(28, 28, 1)),
      tf.keras.layers.MaxPooling2D(),
      tf.keras.layers.Flatten(),
      tf.keras.layers.Dense(64, activation='relu'),
      tf.keras.layers.Dense(10)
  ])
  logits = model(features, training=False)

  if mode == tf.estimator.ModeKeys.PREDICT:
    predictions = {'logits': logits}
    return tf.estimator.EstimatorSpec(labels=labels, predictions=predictions)

  optimizer = tf.compat.v1.train.GradientDescentOptimizer(
      learning_rate=LEARNING_RATE)
  loss = tf.keras.losses.SparseCategoricalCrossentropy(
      from_logits=True, reduction=tf.keras.losses.Reduction.NONE)(labels, logits)
  loss = tf.reduce_sum(loss) * (1. / BATCH_SIZE)
  if mode == tf.estimator.ModeKeys.EVAL:
    return tf.estimator.EstimatorSpec(mode, loss=loss)

  return tf.estimator.EstimatorSpec(
      mode=mode,
      loss=loss,
      train_op=optimizer.minimize(
          loss, tf.compat.v1.train.get_or_create_global_step()))

Catatan: Meskipun kecepatan pembelajaran adalah tetap dalam contoh ini, secara umum mungkin perlu untuk menyesuaikan kecepatan pembelajaran berdasarkan ukuran batch global.

MultiWorkerMirroredStrategy

Untuk melatih model, gunakan instance tf.distribute.experimental.MultiWorkerMirroredStrategy . MultiWorkerMirroredStrategy membuat salinan semua variabel di lapisan model pada setiap perangkat di semua pekerja. Ini menggunakan CollectiveOps , operasi TensorFlow untuk komunikasi kolektif, untuk menggabungkan gradien dan menjaga agar variabel tetap sinkron. Panduan tf.distribute.Strategy memiliki rincian lebih lanjut tentang strategi ini.

strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy()

WARNING:tensorflow:From /tmp/ipykernel_7505/349189047.py:1: _CollectiveAllReduceStrategyExperimental.__init__ (from tensorflow.python.distribute.collective_all_reduce_strategy) is deprecated and will be removed in a future version.
Instructions for updating:
use distribute.MultiWorkerMirroredStrategy instead
INFO:tensorflow:Single-worker MultiWorkerMirroredStrategy with local_devices = ('/device:GPU:0',), communication = CommunicationImplementation.AUTO

Latih dan evaluasi modelnya

Selanjutnya, tentukan strategi distribusi di RunConfig untuk estimator, dan latih dan evaluasi dengan menjalankan tf.estimator.train_and_evaluate . Tutorial ini hanya mendistribusikan pelatihan dengan menentukan strategi melalui train_distribute . Dimungkinkan juga untuk mendistribusikan evaluasi melalui eval_distribute .

config = tf.estimator.RunConfig(train_distribute=strategy)

classifier = tf.estimator.Estimator(
    model_fn=model_fn, model_dir='/tmp/multiworker', config=config)
tf.estimator.train_and_evaluate(
    classifier,
    train_spec=tf.estimator.TrainSpec(input_fn=input_fn),
    eval_spec=tf.estimator.EvalSpec(input_fn=input_fn)
)

INFO:tensorflow:Initializing RunConfig with distribution strategies.
INFO:tensorflow:Not using Distribute Coordinator.
INFO:tensorflow:Using config: {'_model_dir': '/tmp/multiworker', '_tf_random_seed': None, '_save_summary_steps': 100, '_save_checkpoints_steps': None, '_save_checkpoints_secs': 600, '_session_config': allow_soft_placement: true
graph_options {
  rewrite_options {
    meta_optimizer_iterations: ONE
  }
}
, '_keep_checkpoint_max': 5, '_keep_checkpoint_every_n_hours': 10000, '_log_step_count_steps': 100, '_train_distribute': <tensorflow.python.distribute.collective_all_reduce_strategy._CollectiveAllReduceStrategyExperimental object at 0x7f3404234490>, '_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({}), '_task_type': 'worker', '_task_id': 0, '_global_id_in_cluster': 0, '_master': '', '_evaluation_master': '', '_is_chief': True, '_num_ps_replicas': 0, '_num_worker_replicas': 1, '_distribute_coordinator_mode': None}
INFO:tensorflow:Not using Distribute Coordinator.
INFO:tensorflow:Running training and evaluation locally (non-distributed).
INFO:tensorflow:Start train and evaluate loop. The evaluate will happen after every checkpoint. Checkpoint frequency is determined based on RunConfig arguments: save_checkpoints_steps None or save_checkpoints_secs 600.
WARNING:tensorflow:From /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow_estimator/python/estimator/estimator.py:1244: StrategyBase.configure (from tensorflow.python.distribute.distribute_lib) is deprecated and will be removed in a future version.
Instructions for updating:
use `update_config_proto` instead.
INFO:tensorflow:The `input_fn` accepts an `input_context` which will be given by DistributionStrategy
INFO:tensorflow:Calling model_fn.
/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/data/ops/dataset_ops.py:449: UserWarning: To make it possible to preserve tf.data options across serialization boundaries, their implementation has moved to be part of the TensorFlow graph. As a consequence, the options value is in general no longer known at graph construction time. Invoking this method in graph mode retains the legacy behavior of the original implementation, but note that the returned value might not reflect the actual value of the options.
  warnings.warn("To make it possible to preserve tf.data options across "
INFO:tensorflow:Calling model_fn.
INFO:tensorflow:Done calling model_fn.
INFO:tensorflow:Done calling model_fn.
INFO:tensorflow:Create CheckpointSaverHook.
INFO:tensorflow:Create CheckpointSaverHook.
WARNING:tensorflow:From /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow_estimator/python/estimator/util.py:95: DistributedIteratorV1.initialize (from tensorflow.python.distribute.v1.input_lib) is deprecated and will be removed in a future version.
Instructions for updating:
Use the iterator's `initializer` property instead.
WARNING:tensorflow:From /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow_estimator/python/estimator/util.py:95: DistributedIteratorV1.initialize (from tensorflow.python.distribute.v1.input_lib) is deprecated and will be removed in a future version.
Instructions for updating:
Use the iterator's `initializer` property instead.
INFO:tensorflow:Graph was finalized.
INFO:tensorflow:Graph was finalized.
INFO:tensorflow:Running local_init_op.
INFO:tensorflow:Running local_init_op.
INFO:tensorflow:Done running local_init_op.
INFO:tensorflow:Done running local_init_op.
INFO:tensorflow:Calling checkpoint listeners before saving checkpoint 0...
INFO:tensorflow:Calling checkpoint listeners before saving checkpoint 0...
INFO:tensorflow:Saving checkpoints for 0 into /tmp/multiworker/model.ckpt.
INFO:tensorflow:Saving checkpoints for 0 into /tmp/multiworker/model.ckpt.
INFO:tensorflow:Calling checkpoint listeners after saving checkpoint 0...
INFO:tensorflow:Calling checkpoint listeners after saving checkpoint 0...
2022-01-26 05:29:43.503603: W tensorflow/core/grappler/utils/graph_view.cc:836] No registered 'MultiDeviceIteratorFromStringHandle' OpKernel for GPU devices compatible with node { {node MultiDeviceIteratorFromStringHandle} }
    .  Registered:  device='CPU'

2022-01-26 05:29:43.504873: W tensorflow/core/grappler/utils/graph_view.cc:836] No registered 'MultiDeviceIteratorGetNextFromShard' OpKernel for GPU devices compatible with node { {node MultiDeviceIteratorGetNextFromShard} }
    .  Registered:  device='CPU'
INFO:tensorflow:loss = 2.292878, step = 0
INFO:tensorflow:loss = 2.292878, step = 0
INFO:tensorflow:global_step/sec: 173.275
INFO:tensorflow:global_step/sec: 173.275
INFO:tensorflow:loss = 2.29561, step = 100 (0.579 sec)
INFO:tensorflow:loss = 2.29561, step = 100 (0.579 sec)
INFO:tensorflow:global_step/sec: 189.057
INFO:tensorflow:global_step/sec: 189.057
INFO:tensorflow:loss = 2.2644367, step = 200 (0.529 sec)
INFO:tensorflow:loss = 2.2644367, step = 200 (0.529 sec)
INFO:tensorflow:global_step/sec: 193.075
INFO:tensorflow:global_step/sec: 193.075
INFO:tensorflow:loss = 2.2662685, step = 300 (0.517 sec)
INFO:tensorflow:loss = 2.2662685, step = 300 (0.517 sec)
INFO:tensorflow:global_step/sec: 199.957
INFO:tensorflow:global_step/sec: 199.957
INFO:tensorflow:loss = 2.2667098, step = 400 (0.500 sec)
INFO:tensorflow:loss = 2.2667098, step = 400 (0.500 sec)
INFO:tensorflow:global_step/sec: 204.217
INFO:tensorflow:global_step/sec: 204.217
INFO:tensorflow:loss = 2.251912, step = 500 (0.490 sec)
INFO:tensorflow:loss = 2.251912, step = 500 (0.490 sec)
INFO:tensorflow:global_step/sec: 201.747
INFO:tensorflow:global_step/sec: 201.747
INFO:tensorflow:loss = 2.2633677, step = 600 (0.496 sec)
INFO:tensorflow:loss = 2.2633677, step = 600 (0.496 sec)
INFO:tensorflow:global_step/sec: 206.079
INFO:tensorflow:global_step/sec: 206.079
INFO:tensorflow:loss = 2.2531767, step = 700 (0.485 sec)
INFO:tensorflow:loss = 2.2531767, step = 700 (0.485 sec)
INFO:tensorflow:global_step/sec: 231.299
INFO:tensorflow:global_step/sec: 231.299
INFO:tensorflow:loss = 2.2578738, step = 800 (0.433 sec)
INFO:tensorflow:loss = 2.2578738, step = 800 (0.433 sec)
INFO:tensorflow:global_step/sec: 657.044
INFO:tensorflow:global_step/sec: 657.044
INFO:tensorflow:loss = 2.2344787, step = 900 (0.150 sec)
INFO:tensorflow:loss = 2.2344787, step = 900 (0.150 sec)
INFO:tensorflow:Calling checkpoint listeners before saving checkpoint 938...
INFO:tensorflow:Calling checkpoint listeners before saving checkpoint 938...
INFO:tensorflow:Saving checkpoints for 938 into /tmp/multiworker/model.ckpt.
INFO:tensorflow:Saving checkpoints for 938 into /tmp/multiworker/model.ckpt.
INFO:tensorflow:Calling checkpoint listeners after saving checkpoint 938...
INFO:tensorflow:Calling checkpoint listeners after saving checkpoint 938...
INFO:tensorflow:Calling model_fn.
INFO:tensorflow:Calling model_fn.
INFO:tensorflow:Done calling model_fn.
INFO:tensorflow:Done calling model_fn.
INFO:tensorflow:Starting evaluation at 2022-01-26T05:29:56
INFO:tensorflow:Starting evaluation at 2022-01-26T05:29:56
INFO:tensorflow:Graph was finalized.
INFO:tensorflow:Graph was finalized.
INFO:tensorflow:Restoring parameters from /tmp/multiworker/model.ckpt-938
INFO:tensorflow:Restoring parameters from /tmp/multiworker/model.ckpt-938
INFO:tensorflow:Running local_init_op.
INFO:tensorflow:Running local_init_op.
INFO:tensorflow:Done running local_init_op.
INFO:tensorflow:Done running local_init_op.
INFO:tensorflow:Evaluation [10/100]
INFO:tensorflow:Evaluation [10/100]
INFO:tensorflow:Evaluation [20/100]
INFO:tensorflow:Evaluation [20/100]
INFO:tensorflow:Evaluation [30/100]
INFO:tensorflow:Evaluation [30/100]
INFO:tensorflow:Evaluation [40/100]
INFO:tensorflow:Evaluation [40/100]
INFO:tensorflow:Evaluation [50/100]
INFO:tensorflow:Evaluation [50/100]
INFO:tensorflow:Evaluation [60/100]
INFO:tensorflow:Evaluation [60/100]
INFO:tensorflow:Evaluation [70/100]
INFO:tensorflow:Evaluation [70/100]
INFO:tensorflow:Evaluation [80/100]
INFO:tensorflow:Evaluation [80/100]
INFO:tensorflow:Evaluation [90/100]
INFO:tensorflow:Evaluation [90/100]
INFO:tensorflow:Evaluation [100/100]
INFO:tensorflow:Evaluation [100/100]
INFO:tensorflow:Inference Time : 2.04637s
INFO:tensorflow:Inference Time : 2.04637s
INFO:tensorflow:Finished evaluation at 2022-01-26-05:29:58
INFO:tensorflow:Finished evaluation at 2022-01-26-05:29:58
INFO:tensorflow:Saving dict for global step 938: global_step = 938, loss = 2.234131
INFO:tensorflow:Saving dict for global step 938: global_step = 938, loss = 2.234131
INFO:tensorflow:Saving 'checkpoint_path' summary for global step 938: /tmp/multiworker/model.ckpt-938
INFO:tensorflow:Saving 'checkpoint_path' summary for global step 938: /tmp/multiworker/model.ckpt-938
INFO:tensorflow:Loss for final step: 1.10881.
INFO:tensorflow:Loss for final step: 1.10881.
({'loss': 2.234131, 'global_step': 938}, [])

Optimalkan kinerja pelatihan

Anda sekarang memiliki model dan Pengukur yang mampu multi-pekerja yang didukung oleh tf.distribute.Strategy . Anda dapat mencoba teknik berikut untuk mengoptimalkan kinerja pelatihan multi-pekerja:

• Tingkatkan ukuran batch: Ukuran batch yang ditentukan di sini adalah per-GPU. Secara umum, ukuran batch terbesar yang sesuai dengan memori GPU disarankan.
• Keluarkan variabel: Keluarkan variabel ke tf.float jika memungkinkan. Model ResNet resmi menyertakan contoh bagaimana hal ini dapat dilakukan.

• Gunakan komunikasi kolektif: MultiWorkerMirroredStrategy menyediakan beberapa implementasi komunikasi kolektif .

  • RING mengimplementasikan kolektif berbasis cincin menggunakan gRPC sebagai lapisan komunikasi lintas-host.
  • NCCL menggunakan NCCL Nvidia untuk mengimplementasikan kolektif.
  • AUTO menunda pilihan ke runtime.

  Pilihan terbaik untuk implementasi kolektif tergantung pada jumlah dan jenis GPU, dan interkoneksi jaringan dalam cluster. Untuk mengesampingkan pilihan otomatis, tentukan nilai yang valid ke parameter communication konstruktor MultiWorkerMirroredStrategy , misalnya communication=tf.distribute.experimental.CollectiveCommunication.NCCL .

Kunjungi bagian Performa dalam panduan untuk mempelajari lebih lanjut tentang strategi dan fitur lain yang dapat Anda gunakan untuk mengoptimalkan performa model TensorFlow Anda.

Contoh kode lainnya

1. Contoh ujung ke ujung untuk pelatihan multi pekerja di tensorflow/ekosistem menggunakan template Kubernetes. Contoh ini dimulai dengan model Keras dan mengonversinya menjadi Penaksir menggunakan API tf.keras.estimator.model_to_estimator .
2. Model resmi , banyak di antaranya dapat dikonfigurasi untuk menjalankan beberapa strategi distribusi.