使用MirroredStrategy多GPU训练时验证阶段出现Segmentation Fault问题求助

我现在在用TensorFlow 2.18.0结合tf.distribute.MirroredStrategy在两块GPU上训练模型。单GPU训练一切正常，但切换到双GPU后，验证阶段直接触发了segmentation fault（段错误）。

下面是我的代码片段：

from config import MainConfig
from dataset import dataset
from model2 import build_tf_model
from utils import CustomModelCheckpoint, get_lr_callback
import tensorflow as tf

checkpoint_callback_val = CustomModelCheckpoint(
    "models/val_model_{epoch:02d}_{val_acc_l:.1f}.keras",
    monitor="val_acc_l",
    save_best_only=True,
    mode="max",
    verbose=0,
    start_epoch=5
)

gpus = tf.config.experimental.list_physical_devices('GPU')
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)

input_shape = (MainConfig.sequence_length, MainConfig.features)
train_sequences, train_labels, validation_sequences, validation_labels = dataset() # numpy arrays of shape (samples, sequences, features)
strategy = tf.distribute.MirroredStrategy(devices=["/GPU:0", "/GPU:1"])

with strategy.scope():    
    model = build_tf_model(input_shape)
    model.fit(train_sequences, train_labels,
        validation_data=(validation_sequences, validation_labels),
        epochs=MainConfig.epochs,
        shuffle=True,
        batch_size=MainConfig.train_batch_size,
        callbacks=[checkpoint_callback_val, get_lr_callback()]
    )

针对验证数据集，我已经尝试了以下几种方式：

• 直接传入NumPy数组格式的验证数据
• 改用tf.data.Dataset对象作为验证数据集
• 用strategy.experimental_distribute_dataset创建分布式验证数据集

但不管哪种方式，双GPU下都会触发段错误，以下是堆栈追踪信息：

File "/usr/local/lib/python3.11/dist-packages/tensorflow/python/ops/gen_experimental_dataset_ops.py", line 335 in auto_shard_dataset
File "/usr/local/lib/python3.11/dist-packages/tensorflow/python/data/experimental/ops/distribute.py", line 74 in __init__
File "/usr/local/lib/python3.11/dist-packages/tensorflow/python/distribute/input_ops.py", line 56 in auto_shard_dataset
File "/usr/local/lib/python3.11/dist-packages/tensorflow/python/distribute/input_lib.py", line 919 in _create_cloned_datasets_from_dataset
File "/usr/local/lib/python3.11/dist-packages/tensorflow/python/distribute/input_lib.py", line 834 in build
File "/usr/local/lib/python3.11/dist-packages/tensorflow/python/distribute/input_lib.py", line 804 in __init__
File "/usr/local/lib/python3.11/dist-packages/tensorflow/python/distribute/input_util.py", line 65 in get_distributed_dataset
File "/usr/local/lib/python3.11/dist-packages/tensorflow/python/distribute/mirrored_strategy.py", line 592 in _experimental_distribute_dataset
File "/usr/local/lib/python3.11/dist-packages/tensorflow/python/distribute/distribute_lib.py", line 1468 in experimental_distribute_dataset
File "/usr/local/lib/python3.11/dist-packages/keras/src/backend/tensorflow/trainer.py", line 668 in __init__
File "/usr/local/lib/python3.11/dist-packages/keras/src/backend/tensorflow/trainer.py", line 334 in fit
File "/usr/local/lib/python3.11/dist-packages/keras/src/utils/traceback_utils.py", line 117 in error_handler
File "/mnt/train/train2.py", line 28 in <module>

有没有朋友遇到过类似的问题？或者能帮我分析下双GPU场景下为啥会出现这种情况？使用MirroredStrategy多GPU训练时，验证数据集有没有什么特殊的配置要点或者注意事项？

提前谢谢大家的帮助！

备注：内容来源于stack exchange，提问作者TGD