2

I'm using the pre-trained MobileNet from Keras and want to train it using TensorFlows federated learning, but I'm always getting an error that the protocol buffer cannot be serialized since the 2GB limit is exceeded. My inputs are 224x224 RGB images.

Edit: I have a dataset of 1000 Images (500 Images per Client) and I'm using 10 epochs per dataset.

My code for creating and training the Keras model:

def create_compiled_keras_model():
    base_model = MobileNet(weights='imagenet', include_top=False)
    x = base_model.output
    x = GlobalAveragePooling2D()(x)
    x = Dense(1024, activation='relu')(x)
    x = Dense(1024, activation='relu')(x)
    x = Dense(512, activation='relu')(x)
    preds = Dense(47, activation='softmax')(x)

    # specify input & output
    model = Model(inputs=base_model.input, outputs=preds)

    # set the first 20 layers of the network to be non-trainable
    for layer in model.layers[:20]:
        layer.trainable = False
    for layer in model.layers[20:]:
        layer.trainable = True

    def loss_fn(y_true, y_pred):
        return tf.reduce_mean(tf.keras.losses.sparse_categorical_crossentropy(
        y_true, y_pred))

    model.compile(
    loss=loss_fn,
    optimizer=gradient_descent.SGD(learning_rate=0.02),
    metrics=[tf.keras.metrics.SparseCategoricalAccuracy()])

    return model


def model_fn():
    keras_model = create_compiled_keras_model()
    return tff.learning.from_compiled_keras_model(keras_model,  sample_batch)


# @test {"output": "ignore"}
iterative_process =   tff.learning.build_federated_averaging_process(model_fn)
state = iterative_process.initialize()

# @test {"timeout": 600, "output": "ignore"}

state, metrics = iterative_process.next(state, federated_train_data)
print('round  1, metrics={}'.format(metrics))

The error I'm receiving:

Traceback (most recent call last):
  File "/Users/recep/Documents/federated_learning.py", line 128, in <module>
    state, metrics = iterative_process.next(state, federated_train_data)
  File "/Users/recep/anaconda3/envs/python37/lib/python3.6/site-packages/tensorflow_federated/python/core/impl/function_utils.py", line 600, in __call__
    return context.invoke(self, arg)
  File "/Users/recep/anaconda3/envs/python37/lib/python3.6/site-packages/tensorflow_federated/python/core/impl/reference_executor.py", line 700, in invoke
    result = computed_comp.value(computed_arg)
  File "/Users/recep/anaconda3/envs/python37/lib/python3.6/site-packages/tensorflow_federated/python/core/impl/reference_executor.py", line 843, in <lambda>
    return ComputedValue(lambda x: self._compute(comp.result, _wrap(x)),
  File "/Users/recep/anaconda3/envs/python37/lib/python3.6/site-packages/tensorflow_federated/python/core/impl/reference_executor.py", line 746, in _compute
    return self._compute_tuple(comp, context)
  File "/Users/recep/anaconda3/envs/python37/lib/python3.6/site-packages/tensorflow_federated/python/core/impl/reference_executor.py", line 802, in _compute_tuple
    computed_v = self._compute(v, context)
  File "/Users/recep/anaconda3/envs/python37/lib/python3.6/site-packages/tensorflow_federated/python/core/impl/reference_executor.py", line 744, in _compute
    return self._compute_call(comp, context)
  File "/Users/recep/anaconda3/envs/python37/lib/python3.6/site-packages/tensorflow_federated/python/core/impl/reference_executor.py", line 784, in _compute_call
    computed_arg = self._compute(comp.argument, context)
  File "/Users/recep/anaconda3/envs/python37/lib/python3.6/site-packages/tensorflow_federated/python/core/impl/reference_executor.py", line 746, in _compute
    return self._compute_tuple(comp, context)
  File "/Users/recep/anaconda3/envs/python37/lib/python3.6/site-packages/tensorflow_federated/python/core/impl/reference_executor.py", line 802, in _compute_tuple
    computed_v = self._compute(v, context)
  File "/Users/recep/anaconda3/envs/python37/lib/python3.6/site-packages/tensorflow_federated/python/core/impl/reference_executor.py", line 744, in _compute
    return self._compute_call(comp, context)
  File "/Users/recep/anaconda3/envs/python37/lib/python3.6/site-packages/tensorflow_federated/python/core/impl/reference_executor.py", line 784, in _compute_call
    computed_arg = self._compute(comp.argument, context)
  File "/Users/recep/anaconda3/envs/python37/lib/python3.6/site-packages/tensorflow_federated/python/core/impl/reference_executor.py", line 746, in _compute
    return self._compute_tuple(comp, context)
  File "/Users/recep/anaconda3/envs/python37/lib/python3.6/site-packages/tensorflow_federated/python/core/impl/reference_executor.py", line 802, in _compute_tuple
    computed_v = self._compute(v, context)
  File "/Users/recep/anaconda3/envs/python37/lib/python3.6/site-packages/tensorflow_federated/python/core/impl/reference_executor.py", line 744, in _compute
    return self._compute_call(comp, context)
  File "/Users/recep/anaconda3/envs/python37/lib/python3.6/site-packages/tensorflow_federated/python/core/impl/reference_executor.py", line 784, in _compute_call
    computed_arg = self._compute(comp.argument, context)
  File "/Users/recep/anaconda3/envs/python37/lib/python3.6/site-packages/tensorflow_federated/python/core/impl/reference_executor.py", line 746, in _compute
    return self._compute_tuple(comp, context)
  File "/Users/recep/anaconda3/envs/python37/lib/python3.6/site-packages/tensorflow_federated/python/core/impl/reference_executor.py", line 802, in _compute_tuple
    computed_v = self._compute(v, context)
  File "/Users/recep/anaconda3/envs/python37/lib/python3.6/site-packages/tensorflow_federated/python/core/impl/reference_executor.py", line 744, in _compute
    return self._compute_call(comp, context)
  File "/Users/recep/anaconda3/envs/python37/lib/python3.6/site-packages/tensorflow_federated/python/core/impl/reference_executor.py", line 784, in _compute_call
    computed_arg = self._compute(comp.argument, context)
  File "/Users/recep/anaconda3/envs/python37/lib/python3.6/site-packages/tensorflow_federated/python/core/impl/reference_executor.py", line 746, in _compute
    return self._compute_tuple(comp, context)
  File "/Users/recep/anaconda3/envs/python37/lib/python3.6/site-packages/tensorflow_federated/python/core/impl/reference_executor.py", line 802, in _compute_tuple
    computed_v = self._compute(v, context)
  File "/Users/recep/anaconda3/envs/python37/lib/python3.6/site-packages/tensorflow_federated/python/core/impl/reference_executor.py", line 744, in _compute
    return self._compute_call(comp, context)
  File "/Users/recep/anaconda3/envs/python37/lib/python3.6/site-packages/tensorflow_federated/python/core/impl/reference_executor.py", line 784, in _compute_call
    computed_arg = self._compute(comp.argument, context)
  File "/Users/recep/anaconda3/envs/python37/lib/python3.6/site-packages/tensorflow_federated/python/core/impl/reference_executor.py", line 746, in _compute
    return self._compute_tuple(comp, context)
  File "/Users/recep/anaconda3/envs/python37/lib/python3.6/site-packages/tensorflow_federated/python/core/impl/reference_executor.py", line 802, in _compute_tuple
    computed_v = self._compute(v, context)
  File "/Users/recep/anaconda3/envs/python37/lib/python3.6/site-packages/tensorflow_federated/python/core/impl/reference_executor.py", line 744, in _compute
    return self._compute_call(comp, context)
  File "/Users/recep/anaconda3/envs/python37/lib/python3.6/site-packages/tensorflow_federated/python/core/impl/reference_executor.py", line 784, in _compute_call
    computed_arg = self._compute(comp.argument, context)
  File "/Users/recep/anaconda3/envs/python37/lib/python3.6/site-packages/tensorflow_federated/python/core/impl/reference_executor.py", line 746, in _compute
    return self._compute_tuple(comp, context)
  File "/Users/recep/anaconda3/envs/python37/lib/python3.6/site-packages/tensorflow_federated/python/core/impl/reference_executor.py", line 802, in _compute_tuple
    computed_v = self._compute(v, context)
  File "/Users/recep/anaconda3/envs/python37/lib/python3.6/site-packages/tensorflow_federated/python/core/impl/reference_executor.py", line 744, in _compute
    return self._compute_call(comp, context)
  File "/Users/recep/anaconda3/envs/python37/lib/python3.6/site-packages/tensorflow_federated/python/core/impl/reference_executor.py", line 784, in _compute_call
    computed_arg = self._compute(comp.argument, context)
  File "/Users/recep/anaconda3/envs/python37/lib/python3.6/site-packages/tensorflow_federated/python/core/impl/reference_executor.py", line 746, in _compute
    return self._compute_tuple(comp, context)
  File "/Users/recep/anaconda3/envs/python37/lib/python3.6/site-packages/tensorflow_federated/python/core/impl/reference_executor.py", line 802, in _compute_tuple
    computed_v = self._compute(v, context)
  File "/Users/recep/anaconda3/envs/python37/lib/python3.6/site-packages/tensorflow_federated/python/core/impl/reference_executor.py", line 744, in _compute
    return self._compute_call(comp, context)
  File "/Users/recep/anaconda3/envs/python37/lib/python3.6/site-packages/tensorflow_federated/python/core/impl/reference_executor.py", line 791, in _compute_call
    result = computed_fn.value(computed_arg)
  File "/Users/recep/anaconda3/envs/python37/lib/python3.6/site-packages/tensorflow_federated/python/core/impl/reference_executor.py", line 871, in <lambda>
    lambda x: my_method(fit_argument(x, arg_type, context)),
  File "/Users/recep/anaconda3/envs/python37/lib/python3.6/site-packages/tensorflow_federated/python/core/impl/reference_executor.py", line 914, in _federated_map
    fn(ComputedValue(x, mapping_type.parameter)).value for x in arg.value[1]
  File "/Users/recep/anaconda3/envs/python37/lib/python3.6/site-packages/tensorflow_federated/python/core/impl/reference_executor.py", line 914, in <listcomp>
    fn(ComputedValue(x, mapping_type.parameter)).value for x in arg.value[1]
  File "/Users/recep/anaconda3/envs/python37/lib/python3.6/site-packages/tensorflow_federated/python/core/impl/reference_executor.py", line 775, in <lambda>
    return ComputedValue(lambda x: run_tensorflow(comp, x),
  File "/Users/recep/anaconda3/envs/python37/lib/python3.6/site-packages/tensorflow_federated/python/core/impl/reference_executor.py", line 340, in run_tensorflow
    comp.proto, stamped_arg, graph))
  File "/Users/recep/anaconda3/envs/python37/lib/python3.6/site-packages/tensorflow_federated/python/core/impl/tensorflow_deserialization.py", line 90, in deserialize_and_call_tf_computation
    arg_type, arg_binding = graph_utils.capture_result_from_graph(arg, graph)
  File "/Users/recep/anaconda3/envs/python37/lib/python3.6/site-packages/tensorflow_federated/python/core/impl/graph_utils.py", line 244, in capture_result_from_graph
    for k, v in name_value_pairs
  File "/Users/recep/anaconda3/envs/python37/lib/python3.6/site-packages/tensorflow_federated/python/core/impl/graph_utils.py", line 244, in <listcomp>
    for k, v in name_value_pairs
  File "/Users/recep/anaconda3/envs/python37/lib/python3.6/site-packages/tensorflow_federated/python/core/impl/graph_utils.py", line 265, in capture_result_from_graph
    handle_name = result.make_one_shot_iterator().string_handle().name
  File "/Users/recep/anaconda3/envs/python37/lib/python3.6/site-packages/tensorflow/python/data/ops/dataset_ops.py", line 1358, in make_one_shot_iterator
    _make_dataset.add_to_graph(ops.get_default_graph())
  File "/Users/recep/anaconda3/envs/python37/lib/python3.6/site-packages/tensorflow/python/framework/function.py", line 490, in add_to_graph
    self._create_definition_if_needed()
  File "/Users/recep/anaconda3/envs/python37/lib/python3.6/site-packages/tensorflow/python/framework/function.py", line 341, in _create_definition_if_needed
    self._create_definition_if_needed_impl()
  File "/Users/recep/anaconda3/envs/python37/lib/python3.6/site-packages/tensorflow/python/framework/function.py", line 417, in _create_definition_if_needed_impl
    self._op_def = self.definition.signature
  File "/Users/recep/anaconda3/envs/python37/lib/python3.6/site-packages/tensorflow/python/framework/function.py", line 289, in definition
    c_api.TF_FunctionToFunctionDef(self._c_func.func, buf)
tensorflow.python.framework.errors_impl.InvalidArgumentError: Cannot serialize protocol buffer of type tensorflow.FunctionDef as the serialized size (3004868250bytes) would be larger than the limit (2147483647 bytes)
Recep26
  • 175
  • 1
  • 9
  • 1
    This is possible a limitation of the current "reference" execution in TFF, which does some interesting gymnastics with datasets. Things to try: (1) reduce the number of images and epochs per client [long term not likely what is desired], (2) try building the TFF package at head, which has new handling for Datasets [or wait for the next release, which should include this]. There is active work on a perfomant simulation executor, pieces of this should get better over time. – Zachary Garrett May 06 '19 at 16:54
  • Thanks, btw: is there a way to use Tensorboard with TFF? – Recep26 May 07 '19 at 14:41
  • 1
    In some sense, TFF only "wraps" TensorFlow; delegating local computation to the TensorFlow runtime (in the simulation executor). Anything TensorFlow can do, TFF should be able to delegate to it. So having TensorFlow output a tf.Summary for use in TensorBoard should "just work". – Zachary Garrett May 16 '19 at 16:15

1 Answers1

1

In version 0.12.0 TensorFlow Federated started using a new eager-mode executor by default which may improve this significantly; there are reports of rounds completing for training a ResNet model.

Zachary Garrett
  • 2,911
  • 15
  • 23