使用 FilePerUserClientData 耗尽 RAM答案

【问题标题】：Running Out of RAM using FilePerUserClientData使用 FilePerUserClientData 耗尽 RAM
【发布时间】：2022-01-17 02:57:49
【问题描述】：

我在使用 tff.simulation.FilePerUserClientData 进行训练时遇到问题 - 在每轮 10 个客户端进行 5-6 轮后，我很快就会耗尽内存。每一轮的内存使用量都在稳步增加。我试图缩小范围并意识到问题不是实际的迭代过程，而是客户端数据集的创建。简单地在循环中调用create_tf_dataset_for_client(client) 会导致问题。

所以这是我的代码的最小版本：

import tensorflow as tf
import tensorflow_federated as tff
import numpy as np
import pickle

BATCH_SIZE = 16
EPOCHS = 2
MAX_SEQUENCE_LEN = 20
NUM_ROUNDS = 100
CLIENTS_PER_ROUND = 10

def decode_fn(record_bytes):
    return tf.io.parse_single_example(
        record_bytes,
        {"x": tf.io.FixedLenFeature([MAX_SEQUENCE_LEN], dtype=tf.string),
         "y": tf.io.FixedLenFeature([MAX_SEQUENCE_LEN], dtype=tf.string)}
    )

def dataset_fn(path):
    return tf.data.TFRecordDataset([path]).map(decode_fn).padded_batch(BATCH_SIZE).repeat(EPOCHS)

def sample_client_data(data, client_ids, sampling_prob):
    clients_total = len(client_ids)
    x = np.random.uniform(size=clients_total)
    sampled_ids = [client_ids[i] for i in range(clients_total) if x[i] < sampling_prob]
    data = [train_data.create_tf_dataset_for_client(client) for client in sampled_ids]
    return data
    
with open('users.pkl', 'rb') as f:
    users = pickle.load(f)
    
train_client_ids = users["train"]
client_id_to_train_file = {i: "reddit_leaf_tf/" + i for i in train_client_ids}

train_data = tff.simulation.datasets.FilePerUserClientData(
    client_ids_to_files=client_id_to_train_file,
    dataset_fn=dataset_fn
)

sampling_prob = CLIENTS_PER_ROUND / len(train_client_ids)

for round_num in range(0, NUM_ROUNDS):
    print('Round {r}'.format(r=round_num))
    participants_data = sample_client_data(train_data, train_client_ids, sampling_prob)
    print("Round Completed")

我正在使用 tensorflow 联合 19.0。

我创建客户端数据集的方式是否有问题，或者是否以某种方式预期上一轮的 RAM 未被释放？

【问题讨论】：

为了进一步最小化，如果我们跳过sample_client_data 中的随机性，这会重现吗？例如，如果我们只设置sampled_ids = client_ids[:sampling_prob * len(client_ids)]，问题会重现吗？如果我们仅限于单个客户，例如如果只使用train_client_ids = users["train"][:1] 中的第一个客户端来创建一个小的ClientData？在整个 TFF 和内核 TF（尤其是在数据集中）中，有很多地方正在创建可能无法正确释放的缓冲区和缓存。
感谢您的回复！经过一些测试，我意识到当每轮的客户端数量保持不变时，不会出现此问题，因此似乎与github.com/tensorflow/federated/issues/1215中的问题相同。

标签： tensorflow-datasets tensorflow-federated federated-learning

【解决方案1】：

schmana@ 注意到在每轮更改CLIENTS 放置的基数（不同数量的客户端数据集）时会发生这种情况。这会导致缓存归档，如 http://github.com/tensorflow/federated/issues/1215 中所述。

短期内的解决方法是调用：

tff.framework.get_context_stack().current.executor_factory.clean_up_executors()

在每一轮的开始或结束时。

【讨论】：