[tf.data] graduate rejection_resample API from experimental to tf.data.Dataset

[kvignesh1420]

This PR graduates the tf.data.experimental.rejection_resample API into tf.data.Dataset.rejection_resample by making the following changes:

• Adds the deprecation decorator for the experimental API.
• Add the rejection_resample() method to DatasetV2 class.
• Updates example in documentation with new API.
• Regenerate golden API's.
• Moved and updated the rejection_resample_test target from experimental/kernel_tests to kernel_tests
• Updated the RELEASE.md file

TEST LOG

INFO: Build completed successfully, 8120 total actions
//tensorflow/python/data/kernel_tests:rejection_resample_test            PASSED in 5.1s
  Stats over 10 runs: max = 5.1s, min = 2.6s, avg = 3.6s, dev = 0.9s

[kvignesh1420]

In [81]: import numpy as np
    ...: import tensorflow as tf
    ...:
    ...: init_dist = [0.6 , 0.4]
    ...: target_dist = [0.5, 0.5]
    ...: num_classes = len(init_dist)
    ...: num_samples = 10000
    ...: data_np = np.random.choice(num_classes, num_samples, p=init_dist)
    ...: dataset = tf.data.Dataset.from_tensor_slices(data_np)
    ...: vals = defaultdict(int)
    ...: for i in dataset:
    ...:   vals[i.numpy()]+=1
    ...: print("Initial distribution: {}".format(vals))
Initial distribution: defaultdict(<class 'int'>, {1: 4040, 0: 5960})

In [82]: resampler = tf.data.experimental.rejection_resample(
    ...:             class_func=lambda x: x,
    ...:             target_dist=target_dist,
    ...:             initial_dist=init_dist)
    ...:
    ...: dataset = dataset.apply(resampler)
    ...:
    ...: from collections import defaultdict
    ...: vals = defaultdict(int)
    ...: for i in dataset:
    ...:   vals[i[-1].numpy()]+=1
    ...: print("Resampled distribution: {}".format(vals))
Proportion of examples rejected by sampler is high: [0.6][0.6 0.4][0 1]
Proportion of examples rejected by sampler is high: [0.6][0.6 0.4][0 1]
Proportion of examples rejected by sampler is high: [0.6][0.6 0.4][0 1]
Proportion of examples rejected by sampler is high: [0.6][0.6 0.4][0 1]
Proportion of examples rejected by sampler is high: [0.6][0.6 0.4][0 1]
Proportion of examples rejected by sampler is high: [0.6][0.6 0.4][0 1]
Proportion of examples rejected by sampler is high: [0.6][0.6 0.4][0 1]
Proportion of examples rejected by sampler is high: [0.6][0.6 0.4][0 1]
Proportion of examples rejected by sampler is high: [0.6][0.6 0.4][0 1]
Proportion of examples rejected by sampler is high: [0.6][0.6 0.4][0 1]
Resampled distribution: defaultdict(<class 'int'>, {1: 8080, 0: 5960})

cc: @jsimsa I was playing around with the API and observed some weird behavior. In the above example, it seems like elements are being added to the dataset instead of being removed.

[jsimsa]

The dataset produced by rejection_resampling internally uses sample_from_datasets to sample from a) the original dataset and b) a filtered instance of the original dataset (code).

In other words, it is not unexpected that the number of output elements is greater that the cardinality of original input dataset but it is unexpected that the distribution does not match the target distribution. This is related to an issue we have recently fixed for sample_from_datasets, where the transformation would stop respecting the sampling distribution as some of the input datasets to sample from become empty. I will send a follow up CL which will update rejection_resampling implementation to use the new sample_from_dataset argument, which will fix the issue where the distribution of resampled dataset will fail to match the target distribution (in the case where a most or all elements of the resampled dataset are consumed).

@yangustc07 FYI

With my fix patched, the following program (which I adopted from your example):

import collections
import tensorflow as tf

init_dist = [0.5, 0.5]
target_dist = [0.6, 0.4]
dataset = tf.data.Dataset.range(100000)
x = collections.defaultdict(int)
for i in dataset:
  x[i.numpy() % 2] += 1
print("Initial distribution: {}".format(x))

resampler = tf.data.experimental.rejection_resample(
    class_func=lambda x: x % 2,
    target_dist=target_dist,
    initial_dist=init_dist)
dataset = dataset.apply(resampler)

y = collections.defaultdict(int)
for i in dataset:
  cls, _ = i
  y[cls.numpy()] += 1
print("Resampled distribution: {}".format(y))

Produces the following output:

Initial distribution: defaultdict(<class 'int'>, {0: 50000, 1: 50000})
Proportion of examples rejected by sampler is high: [0.5][0.5 0.5][1 0]
Proportion of examples rejected by sampler is high: [0.5][0.5 0.5][1 0]
Proportion of examples rejected by sampler is high: [0.5][0.5 0.5][1 0]
Proportion of examples rejected by sampler is high: [0.5][0.5 0.5][1 0]
Proportion of examples rejected by sampler is high: [0.5][0.5 0.5][1 0]
Proportion of examples rejected by sampler is high: [0.5][0.5 0.5][1 0]
Proportion of examples rejected by sampler is high: [0.5][0.5 0.5][1 0]
Proportion of examples rejected by sampler is high: [0.5][0.5 0.5][1 0]
Proportion of examples rejected by sampler is high: [0.5][0.5 0.5][1 0]
Proportion of examples rejected by sampler is high: [0.5][0.5 0.5][1 0]
Resampled distribution: defaultdict(<class 'int'>, {0: 75089, 1: 50000})

[kvignesh1420]

PR #49009 has been raised as a pre-requisite for the current one.

[gbaned]

@kvignesh1420 Can you please resolve conflicts? Thanks!

[kvignesh1420]

@gbaned the file changes in this PR conflict with the ones in pre-requisite PR's. I will resolve all the conflicts in this PR once the prereqs are merged. Hope it's fine.

[gbaned]

@kvignesh1420 Can you please resolve conflicts? Thanks!

[jsimsa]

@aaudiber could you please take a look? thanks

[aaudiber]

Thanks @kvignesh1420

[aaudiber]

Moving all of our implementations into dataset_ops.py is making dataset_ops.py quite long and hard to navigate. In later PRs we should consider moving the dataset transformations from dataset_ops.py into their own files, similar to what we do for experimental ops. It would also make graduating experimental ops more straightforward. @jsimsa what are your thoughts?

[jsimsa]

This is a valid concern and I think it would be a useful refactor. That said, we should do this in a manner which does not require updating callsites that import dataset_ops. One option would be that we keep dataset_ops.py as a shim to import symbols from the "per transformation" modules.

[kvignesh1420]

@aaudiber @jsimsa we might also want to consider the circular dependencies and file duplications during the promotion/refactor process if we are going to maintain separate files per transformation and use them in dataset_ops.py. Also, if we are using dataset_ops.py as a shim where the API layout is unaffected for the users, we might have to maintain a new file with the actual functionality of the API's in dataset_ops.py so that circular dependencies can be prevented. WDYT?

[aaudiber]

+1 for making dataset_ops a shim. The refactor should have no user-facing impact. We can use LazyLoader in the dataset impl files to handle the circular reference on dataset_ops.py. Keeping dataset_ops readable is much more important than avoiding circular dependencies between dataset implementations and dataset_ops.

[kvignesh1420]

okay sounds good!