Proper way of saving clients' federated datasets

Question

I would like to train two independent TFF models using emnist dataset. Each model should train on a 1000 distinct participants randomly drawn from the dataset.

Code below

emnist_train, emnist_test = tff.simulation.datasets.emnist.load_data()

participants_ids = np.random.choice(a=emnist_train.client_ids, 
                                    size=1000,
                                    replace=False)

federated_dataset = 
        [data_train.create_tf_dataset_for_client(i) for i in participants_ids]

nested_dataset = tf.data.Dataset.from_tensor_slices(federated_dataset)

Trying to save the dataset

tf.data.experimental.save(nested_dataset, 'model_dataset')

the warning below is generated. However, the save is completed.

E tensorflow/core/framework/dataset.cc:89] The Encode() method is not implemented for DatasetVariantWrapper objects.

The problem occurs upon loading the dataset and trying to inspect its contents

dataset = tf.data.experimental.load('model_dataset', 
                      element_spec= 
                      DatasetSpec(collections.OrderedDict([
                         ('label', TensorSpec(shape=(), dtype=tf.int32)),
                         ('pixels', TensorSpec(shape=(28, 28), dtype=tf.float32))]), 
                      TensorShape([])

# verifying elements
for example in dataset:
        print(example)

Error below

tensorflow.python.framework.errors_impl.DataLossError: Unable to parse tensor from stored proto.

Trying other methods such as pickle.dump and np.save, all resulted in error below

tensorflow.python.framework.errors_impl.InternalError: Tensorflow type 21 not convertible to numpy dtype.

Is there any good way to save the newly created datasets ?

Zachary_Garrett_14692 · Answer

Instead of saving a dataset of datasets, would it work to save the client ids that were sampled and construct datasets upon loading?

# Create a dataset of the participating IDs.
id_ds = tf.data.Dataset.from_tensor_slices(participants_ids)
tf.data.experimental.save(id_ds, '/tmp/id_dataset')

# Loaded the dataset back later.
loaded_ds = tf.data.experimental.load(
  '/tmp/id_dataset',
  element_spec=tf.TensorSpec(shape=[], dtype=tf.string))

# Create a federated dataset that yield (client_id, dataset).
federated_dataset = loaded_ds.map(
    lambda id: (id, emnist_train.serializable_dataset_fn(id)))
print(f'Loaded dataset with {tf.data.Dataset.cardinality(federated_dataset)} clients')
>>> Loaded dataset with 1000 clients.

print(f'Dataset element types: {next(iter(federated_dataset))[1].element_spec}')
>>> Dataset element types: OrderedDict([('label', TensorSpec(shape=(), dtype=tf.int32, name=None)), ('pixels', TensorSpec(shape=(28, 28), dtype=tf.float32, name=None))])

for id, dataset in federated_dataset.take(5):
  print(f'Client [{id}] has {sum(1 for _ in iter(dataset))} examples')
>>> Client [b'f2174_61'] has 106 examples
>>> Client [b'f1378_08'] has 99 examples
>>> Client [b'f1550_34'] has 106 examples
>>> Client [b'f3817_22'] has 106 examples
>>> Client [b'f1000_45'] has 109 examples

Then a flattened dataset could be created by replacing tf.data.Dataset.map with tf.data.Dataset.flat_map:

flat_dataset = loaded_ds.flat_map(
    lambda id: emnist_train.serializable_dataset_fn(id))

print(f'Dataset element types: {next(iter(federated_dataset))[1].element_spec}')
>>> Dataset element types: OrderedDict([('label', TensorSpec(shape=(), dtype=tf.int32, name=None)), ('pixels', TensorSpec(shape=(28, 28), dtype=tf.float32, name=None))])

print(f'Flat dataset has {sum(1 for _ in flat_dataset):,} examples.')
>>> Flat dataset has 101,619 examples.

Proper way of saving clients' federated datasets

1 Answers