0
$\begingroup$

I was trying to compute the mean of the images that I fetched from a GCS bucket using TensorFlow's tf.data input pipeline. For this, I came across two methods: tf.math.reduce_mean() and tf.data.Dataset.reduce().

To compare the methods tf.math.reduce_mean() and tf.data.Dataset.reduce(), I took samples of the dataset and noted the time taken by both. I used the following code:

size = [50, 100, 150, 200, 500]
df = pd.DataFrame(
    columns=["Size", "tf.math.reduce_mean()", "tf.data.Dataset.reduce()"]
)

for i in range(len(size)):
    # Getting the sample:
    sample_dataset = train_dataset.take(size[i])

    # Timing tf.math.reduce_mean() method:
    tf_math_start = time.time()     # time starts
    mean_img = tf.math.reduce_mean(
        input_tensor = list(sample_dataset.as_numpy_iterator()),
        axis = 0
    )
    tf_math_end = time.time()       # time stops

    # Timing tf.data.Dataset.reduce() method:
    tf_data_start = time.time()     # time starts
    mean_img = sample_dataset.reduce(
        initial_state = tf.zeros(shape=(256,256), dtype=tf.float64),
        reduce_func = lambda x,y: x+y
    )/size[i]
    tf_data_end = time.time()       # time stops

    # Appending to Dataframe:
    temp_df = pd.DataFrame(
        data = {
            "Size": [size[i]],
            "tf.math.reduce_mean()": [tf_math_end - tf_math_start],
            "tf.data.Dataset.reduce()": [tf_data_end - tf_data_start]
        }
    )
    df = pd.concat([df, temp_df], ignore_index=True)

print(df)

After execution, I got the following results (execution time is in seconds):

  Size  tf.math.reduce_mean()  tf.data.Dataset.reduce()
0   50              51.278539                 57.186869
1  100             106.780465                 81.493149
2  150             120.286897                104.921393
3  200             139.838593                124.364554
4  500             433.217095                320.542538

My questions are:

Ques 1: Does tf.data.Dataset.reduce() call __next__() method internally?

Ques 2: If yes, then why is it still faster than tf.math.reduce()?

Ques 3: If no, then how does tf.data.Dataset.reduce() work, and how can one make its execution faster?

Finally, any other suggestions to speed up the operations are welcome.

Thanks in advance.

P.S.:

The dataset was generated using the following method:

def create_dataset(file_pattern: str) -> tf.data.Dataset:

    dataset = tf.data.Dataset.list_files(file_pattern)

    dataset = dataset.map(
        map_func = tf.io.read_file,
        num_parallel_calls=tf.data.AUTOTUNE
    )

    dataset = dataset.map(
        map_func = tf.image.decode_jpeg,
        num_parallel_calls = tf.data.AUTOTUNE
    )

    dataset = dataset.map(
        map_func = lambda x: tf.image.resize(x,(256,256))[:,:,0]/255,
        num_parallel_calls = tf.data.AUTOTUNE
    )

    dataset = dataset.map(
        map_func = lambda x: tf.cast(x, tf.float64),
        num_parallel_calls = tf.data.AUTOTUNE
    )

    return dataset


train_dataset = create_dataset(
    file_pattern="gs://"+GCS_BUCKET_NAME+"/train/*/*.jpeg"
)
$\endgroup$

0

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.