多GPU TFF仿真错误“在多GPU TFF仿真中检测数据集减少op”

Question

我使用Tensorflow联邦模拟运行了情绪检测模型的代码。我的代码非常好，只使用CPU。但是，当我试图使用GPU运行TFF时，我收到了这个错误。

ValueError: Detected dataset reduce op in multi-GPU TFF simulation: `use_experimental_simulation_loop=True` for `tff.learning`; or use `for ... in iter(dataset)` for your own dataset iteration.Reduce op will be functional after b/159180073.

这个错误是关于什么的，我如何修复它？我试图搜索许多地方，但没有找到答案。

这是调用堆栈，如果有帮助的话。它很长，所以我粘贴到这个链接：https://pastebin.com/b1R93gf1

编辑：

下面是包含iterative_process的代码

def startTraining(output_file):
    
    iterative_process = tff.learning.build_federated_averaging_process(
        model_fn,
        client_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=0.01),
        server_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=1.0),
        use_experimental_simulation_loop=True
    )
    
    flstate = iterative_process.initialize()
    evaluation = tff.learning.build_federated_evaluation(model_fn)
    
    output_file.write(
        'round,available_users,loss,sparse_categorical_accuracy,val_loss,val_sparse_categorical_accuracy,test_loss,test_sparse_categorical_accuracy\n')
    curr_round_result = [0,0,100,0,100,0]
    min_val_loss = 100
    for round in range(1,ROUND_COUNT + 1):
        available_users = fetch_available_users_and_increase_time(ROUND_DURATION_AVERAGE + random.randint(-ROUND_DURATION_VARIATION, ROUND_DURATION_VARIATION + 1))
        if(len(available_users) == 0):
            write_to_file(curr_round_result)
            continue
        train_data = make_federated_data(available_users, 'train')
        flstate, metrics = iterative_process.next(flstate, train_data)
        val_data = make_federated_data(available_users, 'val')
        val_metrics = evaluation(flstate.model, val_data)
        
        curr_round_result[0] = round
        curr_round_result[1] = len(available_users)
        curr_round_result[2] = metrics['train']['loss']
        curr_round_result[3] = metrics['train']['sparse_categorical_accuracy']
        curr_round_result[4] = val_metrics['loss']
        curr_round_result[5] = val_metrics['sparse_categorical_accuracy']
        write_to_file(curr_round_result)

以下是make_federated_data的代码

def make_federated_data(users, dataset_type):
    offset = 0
    if(dataset_type == 'val'):
        offset = train_size
    elif(dataset_type == 'test'):
        offset = train_size + val_size
    
    global LOADED_USER
    for id in users:
        if(id + offset not in LOADED_USER):
            LOADED_USER[id + offset] = getDatasetFromFilePath(filepaths[id + offset])

    return [
        LOADED_USER[id + offset]
        for id in users
    ]

Answer 1

TFF确实支持多GPU，正如错误消息所述，有两件事正在发生：

1. 代码使用的是tff.learning，但是使用了False的默认use_experimental_simulation_loop参数值。对于多个GPU，在使用API(包括True )时，必须将其设置为tff.learning.build_federated_averaging_process。例如，使用：

调用

training_process = tff.learning.build_federated_averaging_process(
  ..., use_experimental_simulation_loop=True)

1. 代码包含某个地方的自定义tf.data.Dataset.reduce(...)调用。必须用在dataset上迭代的Python代码来替换它。例如：

result = dataset.reduce(initial_state=0, reduce_func=lambda s, x: s + x)

变成了

s = 0
for x in iter(dataset):
  s += x

Answer 2

我意识到TFF还没有支持多个GPU。因此，我们需要将GPU的可见数限制在1以下：

os.environ["CUDA_VISIBLE_DEVICES"] = "0"