Skip to content

declearn.model.tensorflow.utils.move_layer_to_device

Create a copy of an input keras layer placed on a given device.

This functions creates a copy of the input layer and of all its weights. It may therefore be costful and should be used sparingly, to move away variables on a device where all further computations are expected to be run.

Parameters:

Name Type Description Default
layer tf_keras.layers.Layer

Keras layer that needs moving to another device.

 required
device Union[tf.config.LogicalDevice, str]

Device where to place the layer's weights.

 required

Returns:

Name Type Description
layer tf_keras.layers.Layer

Copy of the input layer, with its weights backed on device.
Source code in declearn/model/tensorflow/utils/_gpu.py 
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
def move_layer_to_device(
    layer: tf_keras.layers.Layer,
    device: Union[tf.config.LogicalDevice, str],
) -> tf_keras.layers.Layer:
    """Create a copy of an input keras layer placed on a given device.

    This functions creates a copy of the input layer and of all its weights.
    It may therefore be costful and should be used sparingly, to move away
    variables on a device where all further computations are expected to be
    run.

    Parameters
    ----------
    layer: tf.keras.layers.Layer
        Keras layer that needs moving to another device.
    device: tf.config.LogicalDevice or str
        Device where to place the layer's weights.

    Returns
    -------
    layer:
        Copy of the input layer, with its weights backed on `device`.
    """
    config = tf_keras.layers.serialize(layer)
    weights = layer.get_weights() if layer.built else None
    with tf.device(device):
        layer = tf_keras.layers.deserialize(config)
        if weights:
            layer.set_weights(weights)
    return layer