多 GPU TFF 模拟错误"Detected dataset reduce op in multi-GPU TFF simulation"

Question

我使用Tensorflow Federated模拟运行了情绪检测模型的代码。我的代码只使用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

编辑：

这是包含迭代_处理的代码

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_accuracyn')
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，但使用默认的use_experimental_simulation_loop参数值False。对于多个GPU，当使用包括tff.learning.build_federated_averaging_process在内的API时，必须将其设置为True。例如，调用时使用：

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

2. 该代码在某处包含一个自定义tf.data.Dataset.reduce(...)调用。这必须替换为在数据集上迭代的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"