文章目錄

• • 1. 單機(jī)多卡 MirroredStrategy
  • 2. 多機(jī)訓(xùn)練 MultiWorkerMirroredStrategy
  • 3. TPU 張量處理單元

學(xué)習(xí)于：簡單粗暴 TensorFlow 2

1. 單機(jī)多卡 MirroredStrategy

# 分布式訓(xùn)練 import tensorflow as tf import tensorflow_datasets as tfds# 1 單機(jī)多卡 MirroredStrategystrategy = tf.distribute.MirroredStrategy() # 指定設(shè)備 strategy = tf.distribute.MirroredStrategy(devices=['/gpu:0']) # ------------------------------------------------ num_epochs = 5 batch_size_per_replica = 64 learning_rate = 1e-4# 定義策略 strategy = tf.distribute.MirroredStrategy()print("設(shè)備數(shù)量：{}".format(strategy.num_replicas_in_sync)) batch_size = batch_size_per_replica * strategy.num_replicas_in_syncdef resize(img, label): # 處理圖片img = tf.image.resize(img, [224, 224]) / 255.0return img, label# 載入貓狗分類數(shù)據(jù)集 dataset = tfds.load("cats_vs_dogs", split=tfds.Split.TRAIN, as_supervised=True) dataset = dataset.map(resize).shuffle(1024).batch(batch_size)# 使用策略 with strategy.scope():# 模型構(gòu)建代碼放入 with model = tf.keras.applications.MobileNetV2(weights=None, classes=2)model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=learning_rate),loss=tf.keras.losses.SparseCategoricalCrossentropy(),metrics=[tf.keras.metrics.sparse_categorical_accuracy])model.fit(dataset, epochs=num_epochs)

2. 多機(jī)訓(xùn)練 MultiWorkerMirroredStrategy

• 相比上面，多了以下配置
• 'task': {'type': 'worker', 'index': 0} 每臺機(jī)器 index 不一樣

num_workers = 2 os.environ['TF_CONFIG'] = json.dumps({'cluster': {'worker': ["localhost:20000", "localhost:20001"]},'task': {'type': 'worker', 'index': 0} # 每臺機(jī)器的 index 不同 })strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy() batch_size = batch_size_per_replica * num_workers

3. TPU 張量處理單元

可以在 Colab 上運(yùn)行

tpu = tf.distribute.cluster_resolver.TPUClusterResolver() tf.config.experimental_connect_to_cluster(tpu) tf.tpu.experimental.initialize_tpu_system(tpu) strategy = tf.distribute.experimental.TPUStrategy(tpu)

總結(jié)

以上是生活随笔為你收集整理的TensorFlow 2.0 - tf.distribute 分布式训练的全部內(nèi)容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯，歡迎將生活随笔推薦給好友。