问题描述
我正在按照 image classification 教程运行 FL 算法。根据预先定义的参与者数量列表，每轮参与者的数量会有所不同。
number_of_participants_each_round =
[108,113,93,92,114,101,94,107,99,118,111,88,86,96,110,80,84,91,120,109,112,119,97,104,103,89,100,108]
联邦数据在开始训练前经过预处理和批处理。
NUM_EPOCHS = 5
BATCH_SIZE = 20
SHUFFLE_BUFFER = 418
PREFETCH_BUFFER = 10
def preprocess(dataset):
def batch_format_fn(element):
return collections.OrderedDict(
x=tf.reshape(element['pixels'],[-1,784]),y=tf.reshape(element['label'],1]))
return dataset.repeat(NUM_EPOCHS).shuffle(SHUFFLE_BUFFER).batch(
BATCH_SIZE).map(batch_format_fn).prefetch(PREFETCH_BUFFER)
def make_federated_data(client_data,client_ids):
return [preprocess(client_data.create_tf_dataset_for_client(x)) for x in client_ids]
federated_train_data = make_federated_data(data_train,data_train.client_ids)
每轮根据federated_train_data[0:expected_total_clients]从number_of_participants_each_round中随机抽取参与者，然后对iterative_process执行45 rounds。
expected_total_clients = 500
round_nums = 45
for round_num in range(0,round_nums):
sampled_clients =
np.random.choice(a=federated_train_data[0:expected_total_clients],size=number_of_participants_each_round[round_num],replace=False)
state,metrics = iterative_process.next(state,list(sampled_clients))
print('round {:2d},metrics={}'.format(round_num + 1,metrics))
问题是 VRAM 的使用在几轮后迅速爆发，它在第 5.5 GB 轮达到 6~7，并以大约 0.8 GB/round 的速度不断增加，直到训练最终在第 25~26 轮崩溃，此时 VRAM 到达 17 GB，创建了 +4000 个 python 线程。
下面的错误信息
F tensorflow/core/platform/default/env.cc:72] Check Failed: ret == 0 (35 vs. 0)Thread creation via pthread_create() Failed.
### 故障排除 ###
将 number_of_participants_each_round 减少到 20 可以完成训练，但内存消耗仍然巨大且不断增长。
每轮以固定数量的参与者运行相同的代码，整个训练过程中内存消耗很好，总共大约 1.5 ~ 2.0 GB VRAM。
expected_total_clients = 500
fixed_client_size_per_round = 100
round_nums = 45
for round_num in range(0,round_nums):
sampled_clients =
np.random.choice(a=federated_train_data[0:expected_total_clients],size=fixed_client_size_per_round,replace=False)
state,list(sampled_clients))
print('round {:2d},metrics))
额外的细节：
OS: MacOS Mojave,10.14.6
python -V: Python 3.8.5 then downgraded to Python 3.7.9
TF version: 2.4.1
TFF version: 0.18.0
Keras version: 2.4.3
这是正常的内存行为还是 bug？是否有任何重构/提示来优化内存消耗？
解决方法
问题是 TFF 运行时进程的 executor stack 中的错误。
完整的详细信息和错误修复如下