Подмодель Trraining вместо полной модели Tensorflow Federated

#python #tensorflow #tensorflow-federated

Вопрос:

Я пытаюсь изменить пример федеративного TensorFlow. Я хочу создать подмодель из исходной модели и использовать вновь созданную для этапа обучения, а затем отправить веса на сервер, чтобы он обновил исходную модель.

Я знаю, что это не должно было быть сделано внутри client_update , но сервер должен отправить правильную подмодель непосредственно клиенту, но сейчас я предпочитаю это делать.

На данный момент у меня есть 2 проблемы:

  1. Похоже, я не могу создать новую модель внутри client_update функции вот так:
     @tf.function
    def client_update(model, dataset, server_message, client_optimizer):
        """Performans client local training of `model` on `dataset`.
        Args:
          model: A `tff.learning.Model`.
          dataset: A 'tf.data.Dataset'.
          server_message: A `BroadcastMessage` from server.
          client_optimizer: A `tf.keras.optimizers.Optimizer`.
        Returns:
          A 'ClientOutput`.
        """
    
        model_weights = model.weights
    
        import dropout_model
        dropout_model = dropout_model.get_dropoutmodel(model)
    
    
        initial_weights = server_message.model_weights
        tf.nest.map_structure(lambda v, t: v.assign(t), model_weights,
                              initial_weights)
        .....
 

Ошибка вот в чем:

 ValueError: tf.function-decorated function tried to create variables on non-first call.
 

Созданная модель выглядит следующим образом:

     def from_original_to_submodel(only_digits=True):
        """The CNN model used in https://arxiv.org/abs/1602.05629.
        Args:
          only_digits: If True, uses a final layer with 10 outputs, for use with the
            digits only EMNIST dataset. If False, uses 62 outputs for the larger
            dataset.
        Returns:
          An uncompiled `tf.keras.Model`.
        """
        data_format = 'channels_last'
        input_shape = [28, 28, 1]
        max_pool = functools.partial(
            tf.keras.layers.MaxPooling2D,
            pool_size=(2, 2),
            padding='same',
            data_format=data_format)
        conv2d = functools.partial(
            tf.keras.layers.Conv2D,
            kernel_size=5,
            padding='same',
            data_format=data_format,
            activation=tf.nn.relu)
        model = tf.keras.models.Sequential([
            conv2d(filters=32, input_shape=input_shape),
            max_pool(),
            conv2d(filters=64),
            max_pool(),
            tf.keras.layers.Flatten(),
            tf.keras.layers.Dense(410, activation=tf.nn.relu), #20% dropout
            tf.keras.layers.Dense(10 if only_digits else 62),
        ])
        return model
    
    def get_dropoutmodel(model):
        keras_model = from_original_to_submodel(only_digits=False)
        loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
        return tff.learning.from_keras_model(keras_model, loss=loss, input_spec=model.input_spec)
 
  1. Is more like a theorical question. I would like to train a sub model like i said, so i would take the original model weights sent from the server initial_weights and for each layer i would assign a sublist of random weights to the submodel weights. For example, initial_weights for the layer 6 contains 100 elements, my new submodel for the same layer has only 40 elements, i would choose from a random with a seed the 40 elements, doing the training and then send the seed to the server, so that he would choose the same indeces and then update only them. Is that correct? My second version was to create still 100 elements(40 random and 60 equal to 0) but i think this will mess the model performance when aggregating on the server side.

EDIT:

I have modified the client_update_fn function like so:

 @tff.tf_computation(tf_dataset_type, server_message_type)
def client_update_fn(tf_dataset, server_message):
    model = model_fn()
    submodel = submodel_fn()
    client_optimizer = client_optimizer_fn()
    return client_update(model, submodel, tf_dataset, server_message, client_optimizer)
 

Adding a new parameter to the function build_federated_averaging_process like so:

 def build_federated_averaging_process(
        model_fn, submodel_fn,
        server_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=1.0),
        client_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=0.1)):
 

And in the main.py i did this:

 def tff_submodel_fn():
    keras_model = create_submodel_dropout(only_digits=False)
    loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
    return tff.learning.from_keras_model(keras_model, loss=loss, input_spec=train_data.element_type_structure)

iterative_process = simple_fedavg_tff.build_federated_averaging_process(
    tff_model_fn, tff_submodel_fn, server_optimizer_fn, client_optimizer_fn)
 

Now inside the client_update i can use the submodel:

 @tf.function
def client_update(model, submodel, dataset, server_message, client_optimizer):
    """Performans client local training of `model` on `dataset`.
    Args:
      model: A `tff.learning.Model`.
      dataset: A 'tf.data.Dataset'.
      server_message: A `BroadcastMessage` from server.
      client_optimizer: A `tf.keras.optimizers.Optimizer`.
    Returns:
      A 'ClientOutput`.
    """



    model_weights = model.weights
    initial_weights = server_message.model_weights      
    submodel_weights = submodel.weights
    tf.nest.map_structure(lambda v, t: v.assign(t), submodel_weights,
                          initial_weights)
    num_examples = tf.constant(0, dtype=tf.int32)
    loss_sum = tf.constant(0, dtype=tf.float32)

    # Explicit use `iter` for dataset is a trick that makes TFF more robust in
    # GPU simulation and slightly more performant in the unconventional usage
    # of large number of small datasets.
    weights_delta = []
    testing = False
    if not testing:
        for batch in iter(dataset):
            with tf.GradientTape() as tape:
                outputs = model.forward_pass(batch)
            grads = tape.gradient(outputs.loss, submodel_weights.trainable)
            client_optimizer.apply_gradients(zip(grads, submodel_weights.trainable))
            batch_size = tf.shape(batch['x'])[0]
            num_examples  = batch_size
            loss_sum  = outputs.loss * tf.cast(batch_size, tf.float32)

        weights_delta = tf.nest.map_structure(lambda a, b: a - b,
                                              submodel_weights.trainable,
                                              initial_weights.trainable)
    client_weight = tf.cast(num_examples, tf.float32)
    return ClientOutput(weights_delta, client_weight, loss_sum / client_weight)
 

Я получаю эту ошибку:

     ValueError: No gradients provided for any variable: ['conv2d_2/kernel:0', 'conv2d_2/bias:0', 'conv2d_3/kernel:0', 'conv2d_3/bias:0', 'dense_2/kernel:0', 'dense_2/bias:0', 'dense_3/kernel:0', 'dense_3/bias:0'].

Fatal Python error: Segmentation fault

Current thread 0x00007f27af18b740 (most recent call first):
  File "virtual-environment/lib/python3.8/site-packages/tensorflow/python/framework/ops.py", line 1853 in _create_c_op
  File "virtual-environment/lib/python3.8/site-packages/tensorflow/python/framework/ops.py", line 2041 in __init__
  File "virtual-environment/lib/python3.8/site-packages/tensorflow/python/framework/ops.py", line 3557 in _create_op_internal
  File "virtual-environment/lib/python3.8/site-packages/tensorflow/python/framework/func_graph.py", line 599 in _create_op_internal
  File "virtual-environment/lib/python3.8/site-packages/tensorflow/python/framework/op_def_library.py", line 748 in _apply_op_helper
  File "virtual-environment/lib/python3.8/site-packages/tensorflow/python/ops/gen_dataset_ops.py", line 1276 in delete_iterator
  File "virtual-environment/lib/python3.8/site-packages/tensorflow/python/data/ops/iterator_ops.py", line 549 in __del__

Process finished with exit code 11
 

На данный момент модель такая же, как и оригинальная, я скопировал функцию create_original_fedavg_cnn_model внутри create_submodel_dropout , поэтому не понимаю, в чем дело

Ответ №1:

В общем случае мы не можем создавать переменные внутри a tf.function , так как метод будет повторно использоваться повторно при вычислении TFF; хотя технически переменные могут быть созданы только один раз внутри a tf.function . Мы видим, что на самом model деле он создан вне tf.function большей части кода библиотеки TFF и передан в качестве аргумента a tf.function (пример: https://github.com/tensorflow/federated/blob/44d012f690005ecf9217e3be970a4f8a356e88ed/tensorflow_federated/python/examples/simple_fedavg/simple_fedavg_tff.py#L101). Еще одной возможностью для изучения может быть tf.init_scope контекст, но обязательно полностью прочитайте всю документацию об предостережениях и поведении.

У TFF есть новый коммуникационный примат под названием tff.federated_select , который может оказаться здесь очень полезным. Встроенный поставляется с двумя учебными пособиями:

  1. Отправка Различных Данных Конкретным Клиентам, С tff.federated_select которыми конкретно обсуждается примитив связи.
  2. Эффективное для клиента федеративное обучение большой модели с помощью federated_select и разреженной агрегации, которое демонстрирует использование federated_select федеративного обучения для линейной регрессии; и демонстрирует необходимость «разреженной агрегации», сложность, которую вы определили с заполнением нулей.

Комментарии:

1. Спасибо за ответ. Не могли бы вы, пожалуйста, взглянуть на редактирование вопроса? Я сделал то, что вы предложили, но я не могу понять, что не так