赞
踩
分布式训练是一种用于在多个设备
或机器
上同时训练深度学习模型的技术,它有助于减少训练时间
,允许使用更多数据更快训练
大模型。分布式训练重点关注数据并行性
,本次试验使用的是单机多卡
的分布式训练策略,也就是 MirroredStrategy
。通常单台机器上的有 1-8 个 GPU , 这也是研究人员和小公司最常见的配置。
tf.distribute.MirroredStrategy 的步骤如下:
模型备份
;分成 N 份
,分别传入 N 个计算设备(即数据并行);梯度
;All-reduce
操作,在所有 GPU 间高效交换梯度数据并进行求和
,使得最终每个设备都有了所有设备的梯度之和;同步
的,所以只有当所有设备均更新模型后,才进入下一轮训练。4090
,所以结果肯定是包含只有一个物理 GPU 的列表,并将它们存储在 physical_devices
列表中。physical_devices[0]
划分成了四个虚拟 GPU 设备,每个虚拟 GPU 的内存限制被设置为 2048MB
。这样是为了模拟一个单机多卡
的分布式环境,方便我们试验 MirroredStrategy
策略。physical_devices = tf.config.list_physical_devices('GPU')
tf.config.set_logical_device_configuration(
physical_devices[0],
[
tf.config.LogicalDeviceConfiguration(memory_limit=2048),
tf.config.LogicalDeviceConfiguration(memory_limit=2048),
tf.config.LogicalDeviceConfiguration(memory_limit=2048),
tf.config.LogicalDeviceConfiguration(memory_limit=2048),
]
)
logical_devices = tf.config.list_logical_devices('GPU')
print(f'从一张物理 GPU 中虚拟出 {len(logical_devices)} 个逻辑 GPU')
结果打印:
从一张物理 GPU 中虚拟出 4 个逻辑 GPU
这里主要是准备用于训练神经网络的文本数据集,并对数据进行一些预处理,具体如下:
128
,训练的总轮数为 5
。 从指定的 URL 下载 wiki 文本数据集,并将其解压缩到本地。train_ds
、val_ds
和 test_ds
:这三个变量分别用于表示训练、验证和测试数据集。对每个数据集进行了类似的处理步骤:
洗牌
,以打乱样本的顺序,有助于模型的训练。BATCH_SIZE
。缓存
,以提高数据加载的效率。预取策略
,允许在模型训练时异步加载下一个批次的数据,以减少训练时的等待取数据的时间。BATCH_SIZE = 128
EPOCHS = 5
keras.utils.get_file(origin="https://s3.amazonaws.com/research.metamind.io/wikitext/wikitext-2-v1.zip", extract=True, )
wiki_dir = os.path.expanduser("~/.keras/datasets/wikitext-2/")
train_ds = (tf.data.TextLineDataset(wiki_dir + 'wiki.train.tokens').filter(lambda x: tf.strings.length(x) > 100).shuffle(buffer_size=500).batch(BATCH_SIZE).cache().prefetch(tf.data.AUTOTUNE))
val_ds = (tf.data.TextLineDataset(wiki_dir + 'wiki.valid.tokens').filter(lambda x: tf.strings.length(x) > 100).shuffle(buffer_size=500).batch(BATCH_SIZE).cache().prefetch(tf.data.AUTOTUNE))
test_ds = (tf.data.TextLineDataset(wiki_dir + 'wiki.test.tokens').filter(lambda x: tf.strings.length(x) > 100).shuffle(buffer_size=500).batch(BATCH_SIZE).cache().prefetch(tf.data.AUTOTUNE))
这里介绍训练一个 BERT 掩码语言模型(Masked Language Model,MLM),并使用分布式训练策略 tf.distribute.MirroredStrategy
进行训练。具体如下:
TRAINING_STEP
定义了记录了每个 epoch 总的训练步骤。lr_schedule
定义了一个学习率衰减策略,学习率在训练过程中从初始值(0.0001)线性地衰减到结束值(0.0),衰减的步数由 TRAINING_STEP
决定,也就是每个 epoch 过后学习率进行衰减下调。callbacks
:定义一些回调函数,用于在训练过程中执行特定的操作。包括了早停(EarlyStopping
)和记录训练日志(TensorBoard
)。BERT
掩码语言模型,其中包括预训练的 BERT 模型,并将其最后一层池化层
设置为不可训练。稀疏分类交叉熵
、优化器为 Adam
和评估指标为稀疏分类准确率
。TRAINING_STEP = sum(1 for _ in train_ds.as_numpy_iterator()) * EPOCHS lr_schedule = tf.keras.optimizers.schedules.PolynomialDecay(initial_learning_rate=0.0001, decay_steps=TRAINING_STEP, end_learning_rate=0.) callbacks = [ tf.keras.callbacks.EarlyStopping(patience=2), tf.keras.callbacks.TensorBoard('./logs')] strategy = tf.distribute.MirroredStrategy() print(f"可用的分布式训练的 GPU 设备有 {strategy.num_replicas_in_sync} 个") with strategy.scope(): model_d = keras_nlp.models.BertMaskedLM.from_preset("bert_tiny_en_uncased") model_d.get_layer("bert_backbone").get_layer("pooled_dense").trainable = False model_d.compile(loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True), optimizer=tf.keras.optimizers.Adam(lr_schedule), weighted_metrics=tf.keras.metrics.SparseCategoricalAccuracy()) model_d.fit(train_ds, validation_data=val_ds, epochs=EPOCHS, callbacks=callbacks, verbose=1) model_d.evaluate(test_ds, verbose=1)
结果打印,理论上训练时间是与显卡数量成反比,卡越多训练时间越快,但是在小数据集中效果不是很明显,因为多块显卡之间的通信、复制、同步信息都会耗时,在实际训练中还要保证既要跑满每张显卡又不能 OOM :
可用的分布式训练的 GPU 设备有 4 个
Epoch 1/5
120/120 [==============================] - 48s 270ms/step - loss: 1.9297 - sparse_categorical_accuracy: 0.0579 - val_loss: 1.7024 - val_sparse_categorical_accuracy: 0.1913
Epoch 2/5
120/120 [==============================] - 29s 241ms/step - loss: 1.6517 - sparse_categorical_accuracy: 0.1644 - val_loss: 1.4522 - val_sparse_categorical_accuracy: 0.2798
Epoch 3/5
120/120 [==============================] - 29s 240ms/step - loss: 1.5088 - sparse_categorical_accuracy: 0.2163 - val_loss: 1.3278 - val_sparse_categorical_accuracy: 0.3198
Epoch 4/5
120/120 [==============================] - 29s 240ms/step - loss: 1.4406 - sparse_categorical_accuracy: 0.2370 - val_loss: 1.2749 - val_sparse_categorical_accuracy: 0.3361
Epoch 5/5
120/120 [==============================] - 29s 241ms/step - loss: 1.4113 - sparse_categorical_accuracy: 0.2448 - val_loss: 1.2603 - val_sparse_categorical_accuracy: 0.3402
15/15 [==============================] - 4s 112ms/step - loss: 1.2633 - sparse_categorical_accuracy: 0.3531
可以看出损失在下降,准确率在提升,可以使用更大的 epoch 继续进行训练。使用 tensorboard 查看训练过程 loos 的变化过程如下:
作为一名热心肠的互联网老兵,我意识到有很多经验和知识值得分享给大家,也可以通过我们的能力和经验解答大家在人工智能学习中的很多困惑,所以在工作繁忙的情况下还是坚持各种整理和分享。
但苦于知识传播途径有限,很多互联网行业朋友无法获得正确的资料得到学习提升,故此将并将重要的 AI大模型资料
包括AI大模型入门学习思维导图、精品AI大模型学习书籍手册、视频教程、实战学习等录播视频免费分享出来。
Copyright © 2003-2013 www.wpsshop.cn 版权所有,并保留所有权利。