Using custom strategy to distribute different features to different devices

Research & Models
, ,
#1

Hey I was trying to create a custom strategy that can distribute some of the features from a tensorflow dataset to CPU while the rest to GPUs. Like the following dataset:
{
“dense_features”: {Replica0_gpu: […], Replica1_gpu: […], Replica2_gpu: […], Replica3_gpu: […]},
“sparse_feature_1”: {Replica0_cpu: […], Replica1_cpu: […], Replica2_cpu: […], Replica3_cpu: […]},
“sparse_feature_2”: …,
…
}
note that i also want to preserve the correct row order between replicas, does anyone have any idea? Thanks!

#3

Hi @bangxi_xiao ,

Below are my inputs:

You can use the tf.data.experimental.AutoShardPolicy class to create a custom strategy that distributes some of the features from a TensorFlow dataset to CPU while the rest to GPUs.

The AutoShardPolicy class takes a dictionary as input, where the keys are the names of the features and the values are the desired distribution of the features. In your case, you would specify that the dense_features should be distributed to GPUs and the sparse_features should be distributed to CPUs.

The AutoShardPolicy class will then automatically shard the dataset according to the specified distribution. It will also preserve the correct row order between replicas.


from tensorflow.data.experimental import AutoShardPolicy

# Create the policy
policy = AutoShardPolicy({
    "dense_features": ["Replica0_gpu", "Replica1_gpu", "Replica2_gpu", "Replica3_gpu"],
    "sparse_features": ["Replica0_cpu", "Replica1_cpu", "Replica2_cpu", "Replica3_cpu"]
})

# Apply the policy
dataset = dataset.apply(policy)

For more detailed explanation you can use following tf.data.experimental.AutoShardPolicy from TF documentation.

I hope this helps you.