Repository: tensorflow
Issue: Prefetch to GPU: prefetch_to_device does not do anything (from_generator), tf.data.Dataset API unclear
Labels: comp:data,type:performance

**System information**
- Have I written custom code (as opposed to using a stock example script provided in TensorFlow): Yes
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Ubuntu 20.04
- TensorFlow installed from (source or binary): Official 2.4-nightly docker (version 2.4.0-dev20201005)
- TensorFlow version (use command below): Official 2.4-nightly docker (version 2.4.0-dev20201005)
- Python version: 3.6.9
- CUDA/cuDNN version: CUDA: 11.1, cuDNN: 8.0.2
- GPU model and memory: Nvidia GTX 1080 Ti, 11170MiB

The API concerning moving CPU to GPU with prefetching extremely unclear using tf.data.Dataset. The function 'prefetch_to_device' simply does not work, even though it was stated that it should be fixed by TF 2.3 or TF 2.4 in the following issue:
[issue 35563](https://github.com/tensorflow/tensorflow/issues/35563)

In order to show the behavior, I have written a standalone test that goes over four options:
 - 0: only use 'prefetch_to_device'
 - 1: Use 'copy_to_device'
 - 2: Use 'copy_to_device' and prefetch
 - 3: Use 'copy_to_device' and 'prefetch_to_device'

I would expect options 0, 2, and 3 to be almost identical. However, only option 2 does actually the indented behavior. (Option 1 does not do prefetching, so in that sense it also does what it is intended for).

I am also not even 100% sure if option 2 is optimal. It does seem that the other GPU ops are blocked during the MemCpy. It could be that it only looks as if it is covered. But I do not have the expertise to judge this case. In terms of timings I did notice that option 2 was the fastest: 

~~~
Option prefetch_gpu
data/logs/pa_20201009-082628_prefetch_gpu
Time lapsed= 0:00:01.177912

Option copy
data/logs/pa_20201009-082628_copy
Time lapsed= 0:00:00.916617

Option copy_prefetch
data/logs/pa_20201009-082628_copy_prefetch
Time lapsed= 0:00:00.831833

Option copy_prefetch_gpu
data/logs/pa_20201009-082628_copy_prefetch_gpu
Time lapsed= 0:00:00.926764
~~~

**Standalone code to reproduce the issue**
```python
import os
import datetime
from tqdm import tqdm

import numpy as np

import tensorflow as tf
print('TF version', tf.__version__)


@tf.function
def do_stuff(wmat, tf_var):

    A = tf.matmul(wmat + 1, tf.transpose(wmat))
    error = tf.reduce_mean(tf_var)
    return error, A 

exp_uuid = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")

n_batches = 20

weights = [None] * n_batches
for i in range(n_batches):
    weights[i] = tf.constant(np.random.rand(2000,10000), dtype=tf.float32)


def gen():
    for i in weights:
        yield i

option_names = ['prefetch_gpu', 'copy', 'copy_prefetch', 'copy_prefetch_gpu']
for option in range(4):

    dataset = tf.data.Dataset.from_generator(gen, output_types=(tf.float32))

    if option == 0:
        ## Option 1: prefetch_gpu
        #
        ## output:  
        ##          weights device /job:localhost/replica:0/task:0/device:CPU:0
        ##          weights device after identity /job:localhost/replica:0/task:0/device:GPU:0

        gpu_transform = tf.data.experimental.prefetch_to_device('/gpu:0')
        dataset.apply(gpu_transform)

    elif option == 1:
        ## Option 1: only copy
        #
        ## output:
        ##          weights device /job:localhost/replica:0/task:0/device:GPU:0
        dataset = dataset.apply(tf.data.experimental.copy_to_device("/gpu:0"))

    elif option == 2:
        ## Option 2: copy + prefetch
        ## as suggested in https://github.com/tensorflow/tensorflow/issues/35563#issuecomment-602160568
        #
        ## output:
        ##          weights device /job:localhost/replica:0/task:0/device:GPU:0

        dataset = dataset.apply(tf.data.experimental.copy_to_device("/gpu:0"))
        with tf.device("/gpu:0"):
            dataset = dataset.prefetch(1)

    elif option == 3:
        ## Option 3: copy + prefetch_gpu
        #
        ## output:
        ##          weights device /job:localhost/replica:0/task:0/device:GPU:0
        dataset = dataset.apply(tf.data.experimental.copy_to_device("/gpu:0"))
        gpu_transform = tf.data.experimental.prefetch_to_device('/gpu:0')
        dataset.apply(gpu_transform)


    tf_var = tf.Variable(np.zeros(3))
    adam = tf.keras.optimizers.Adam(1e-4) 
    logpath = os.path.join('data', 'logs', 'pa_' + exp_uuid + '_' + option_names[option])

    tf.profiler.experimental.start(logpath)
    start = datetime.datetime.now()
    for b, wmat in tqdm(enumerate(dataset)):
        with tf.GradientTape() as tape:

            if b == 0:
                print('\n weights device', wmat.device)
                print('')

            if option == 0:
                wmat = tf.identity(wmat, 'move_to_gpu')
                if b == 0:
                    print('weights device after identity', wmat.device)
                    print('')

            # Do some calculations
            result = do_stuff(wmat, tf_var)
        
        grads = tape.gradient(result[0], [tf_var])
        adam.apply_gradients(zip(grads, [tf_var]))
    stop = datetime.datetime.now()
    tf.profiler.experimental.stop()

    print(f'\nOption {option_names[option]}')
    print(logpath)
    print('Time lapsed=', stop - start)
```

Option 0: prefetch_to_device 
![prefetch_gpu_overview](https://user-images.githubusercontent.com/28840027/95562130-0984eb00-0a1c-11eb-89c0-0949ed49eeb8.png)
![prefetch_gpu_trace](https://user-images.githubusercontent.com/28840027/95562131-0a1d8180-0a1c-11eb-8840-b49acebe77ad.png)

Option 1: copy_to_device
![copy_overview](https://user-images.githubusercontent.com/28840027/95562121-07229100-0a1c-11eb-847e-ce5f0b0dd9e1.png)
![copy_trace](https://user-images.githubusercontent.com/28840027/95562129-08ec5480-0a1c-11eb-8eb9-fc4646f2ee3d.png)

Option 2: copy_to_device + prefetch
![copy_prefetch_overview](https://user-images.githubusercontent.com/28840027/95562125-0853be00-0a1c-11eb-9402-90584ba2360c.png)
![copy_prefetch_trace](https://user-images.githubusercontent.com/28840027/95562127-0853be00-0a1c-11eb-963e-4ef90664a366.png)

Option 3: copy_to_device + prefetch_to_device
![copy_prefetch_gpu_overview](https://user-images.githubusercontent.com/28840027/95562123-07bb2780-0a1c-11eb-947f-27aa82348e1f.png)
![copy_prefetch_gpu_trace](https://user-images.githubusercontent.com/28840027/95562124-07bb2780-0a1c-11eb-9254-47fdc7ab2e81.png)


Files:
[pa_20201009-082628_prefetch_gpu.zip](https://github.com/tensorflow/tensorflow/files/5353593/pa_20201009-082628_prefetch_gpu.zip)
[pa_20201009-082628_copy.zip](https://github.com/tensorflow/tensorflow/files/5353594/pa_20201009-082628_copy.zip)
[pa_20201009-082628_copy_prefetch.zip](https://github.com/tensorflow/tensorflow/files/5353596/pa_20201009-082628_copy_prefetch.zip)
[pa_20201009-082628_copy_prefetch_gpu.zip](https://github.com/tensorflow/tensorflow/files/5353597/pa_20201009-082628_copy_prefetch_gpu.zip)