Exploding memory while training federated model on FLAIR Dataset about federated

[zcharles8]

Thank you for the details! I have a hypothesis about this, looking into it now.

Just to clarify: The memory blowup happens if you do something like

sampled_clients = ...  # Fix a set of clients
for _ in range(100):
  client_data = [train_set.create_tf_dataset_for_client(x) for x in sampled_clients]
  # Do some model training with client_data

Does it happen if you fix the client datasets across all rounds, e.g.:

sampled_clients = ...  # Fix a set of clients
client_data = [train_set.create_tf_dataset_for_client(x) for x in sampled_clients]
for _ in range(100):
  # Do some model training with client_data

[zcharles8]

@christymarc I believe that 1ab31a2 should resolve your issue.

If you'd like to use it with an earlier TFF version, you should be able to do something like the following:

First, sub-class FilePerUserClientData to avoid the issue fixed in 1ab31a2:

class StreamlinedFilePerUserClientData(tff.simulation.datasets.FilePerUserClientData):

  def __init__(
      self,
      client_ids_to_files: Mapping[str, str],
      dataset_fn: Callable[[str], tf.data.Dataset],
  ):
    self._client_ids_to_files = client_ids_to_files
    self._dataset_fn = dataset_fn
    super().__init__(client_ids_to_files, dataset_fn)

  def create_tf_dataset_for_client(self, client_id: str) -> tf.data.Dataset:
    if client_id not in self.client_ids:
      raise ValueError(
          'ID [{i}] is not a client in this ClientData. See '
          'property `client_ids` for the list of valid ids.'.format(i=client_id)
      )
    return self._dataset_fn(self._client_ids_to_files[client_id])

Then, to load flair, do the following as usual:

tff.simulation.datasets.flair.download_data(data_dir)
tff.simulation.datasets.flair.download_data(data_dir, cache_dir)

Then follow the load_data code here but replace FilePerUserClientData with StreamlinedFilePerUserClientData from above.

Please let me know if that solves your problem.

[zcharles8]

@christymarc All tests I've run so far suggest that this has been fixed by the aforementioned change. I'm marking this as resolved for now, but if you continue to hit this issue even with the change (or workaround above), please feel free to re-open.