Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

More rigorous shape inference in to_tf_dataset #4763

Merged
merged 5 commits into from Sep 8, 2022
Merged

More rigorous shape inference in to_tf_dataset #4763

Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
002387e
More rigorous shape inference in to_tf_dataset
Rocketknight1 Jul 28, 2022
a4ef784
Simplify the new shape inference
Rocketknight1 Jul 29, 2022
144304b
Read length from Sequence features instead of just sampling batches
Rocketknight1 Sep 7, 2022
31a6d58
make style
Rocketknight1 Sep 7, 2022
c1b98ee
Remove Sequence-specific code
Rocketknight1 Sep 8, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
8 changes: 3 additions & 5 deletions src/datasets/arrow_dataset.py
View file
Open in desktop
Expand Up @@ -224,7 +224,7 @@ def _get_output_signature(
collate_fn_args: dict,
cols_to_retain: Optional[List[str]] = None,
batch_size: Optional[int] = None,
num_test_batches: int = 10,
num_test_batches: int = 200,
):
"""Private method used by `to_tf_dataset()` to find the shapes and dtypes of samples from this dataset
after being passed through the collate_fn. Tensorflow needs an exact signature for tf.numpy_function, so
Expand Down Expand Up @@ -253,11 +253,9 @@ def _get_output_signature(

if len(dataset) == 0:
raise ValueError("Unable to get the output signature because the dataset is empty.")
if batch_size is None:
test_batch_size = min(len(dataset), 8)
else:
if batch_size is not None:
batch_size = min(len(dataset), batch_size)
test_batch_size = batch_size
test_batch_size = min(len(dataset), 2)

test_batches = []
for _ in range(num_test_batches):
Expand Down