Skip to content

Support empty batches #530

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

Closed
wants to merge 11 commits into from
Closed

Support empty batches #530

Changes from 1 commit
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
b8a2ff7
support empty batches in memory manager and optimizer
Oct 25, 2022
2e1b9d7
restore warning
Oct 25, 2022
df9d1ab
disable functorch test for 1.13+
Oct 25, 2022
0268fa1
Merge branch 'main' of github.com:pytorch/opacus into ffuuugor_522
Oct 26, 2022
b952c2a
0-batch tests
Oct 27, 2022
5c7fc6f
lint
Oct 27, 2022
64f08ad
EW test fix
Oct 27, 2022
df7c355
docstring up
Oct 27, 2022
338097c
typing improvement
Oct 31, 2022
46922c9
merge
Nov 4, 2022
8b11967
merge3
Nov 7, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
typing improvement
  • Loading branch information
Igor Shilov committed Oct 31, 2022
commit 338097cf99e0e970e90b654985df60422b071c20
19 changes: 12 additions & 7 deletions opacus/data_loader.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
# limitations under the License.

import logging
from typing import Any, Optional, Sequence
from typing import Any, Optional, Sequence, Tuple, Type, Union

import torch
from opacus.utils.uniform_sampler import (
Expand All @@ -29,9 +29,10 @@


def wrap_collate_with_empty(
*,
collate_fn: Optional[_collate_fn_t],
Copy link
Contributor

@karthikprasad karthikprasad Oct 27, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
collate_fn: Optional[_collate_fn_t],
*,
collate_fn: Optional[_collate_fn_t],

sample_empty_shapes: Sequence[torch.Size],
dtypes: Sequence[torch.dtype],
sample_empty_shapes: Sequence[Tuple],
dtypes: Sequence[Union[torch.dtype, Type]],
):
"""
Wraps given collate function to handle empty batches.
Expand Down Expand Up @@ -59,7 +60,7 @@ def collate(batch):
return collate


def shape_safe(x: Any):
def shape_safe(x: Any) -> Tuple:
"""
Exception-safe getter for ``shape`` attribute

Expand All @@ -72,7 +73,7 @@ def shape_safe(x: Any):
return x.shape if hasattr(x, "shape") else ()


def dtype_safe(x: Any):
def dtype_safe(x: Any) -> Union[torch.dtype, Type]:
"""
Exception-safe getter for ``dtype`` attribute

Expand Down Expand Up @@ -161,7 +162,7 @@ def __init__(
sample_rate=sample_rate,
generator=generator,
)
sample_empty_shapes = [[0, *shape_safe(x)] for x in dataset[0]]
sample_empty_shapes = [(0, *shape_safe(x)) for x in dataset[0]]
dtypes = [dtype_safe(x) for x in dataset[0]]
Copy link
Contributor

@alexandresablayrolles alexandresablayrolles Oct 26, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If this is consumed by wrap_collate_with_empty, we need the dtype to be an actual dtype and not a type right?

Copy link
Contributor Author

@ffuuugor ffuuugor Oct 27, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

torch.zeros support normal python types (int/float/bool) just fine:

> torch.zeros((2,2), dtype=int)
tensor([[0, 0],
        [0, 0]])

> torch.zeros((2,2), dtype=float)  
tensor([[0., 0.],
        [0., 0.]], dtype=torch.float64)

> torch.zeros((2,2), dtype=bool)   
tensor([[False, False],
        [False, False]])

Copy link

@joserapa98 joserapa98 Oct 27, 2022 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hi there! I've been following the discussion since I ran into the empty batches problem when using poisson sampling with small sampling rate. When I tested this implementation with my own project, it didn't work for me, since it does not support more complex labels. If labels are just numbers (int, float, bool), it's ok, but if you have maybe a label formed by a tuple of numbers, its type will be tuple, thus causing an error in torch.zeros.

I don't know if you expect this kind of things to be supported, but as they work in standard PyTorch (and in fact in Opacus when batches are not empty), maybe it is something worth to be aware of.

Copy link
Contributor Author

@ffuuugor ffuuugor Oct 27, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That's an interesting point, thanks

Just thinking out loud, what would be a good way to support this? In case the original label is a tuple, how does collate function handles it - would it output multiple tensors per label?

Copy link

@joserapa98 joserapa98 Oct 27, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If you have a dataset in which each sample is of type, say, tuple(torch.Tensor, tuple(int, int)), the dataloader would return tuple(torch.Tensor, tuple(torch.Tensor, torch.Tensor)), where now each tensor has an extra dimension for the batch. Something similar would happen if labels are given as lists, dicts, etc.

This is the code snippet that manages this cases:

if isinstance(elem, collections.abc.Mapping):
    try:
        return elem_type({key: collate([d[key] for d in batch], collate_fn_map=collate_fn_map) for key in elem})
    except TypeError:
        # The mapping type may not support `__init__(iterable)`.
        return {key: collate([d[key] for d in batch], collate_fn_map=collate_fn_map) for key in elem}
elif isinstance(elem, tuple) and hasattr(elem, '_fields'):  # namedtuple
    return elem_type(*(collate(samples, collate_fn_map=collate_fn_map) for samples in zip(*batch)))
elif isinstance(elem, collections.abc.Sequence):
    # check to make sure that the elements in batch have consistent size
    it = iter(batch)
    elem_size = len(next(it))
    if not all(len(elem) == elem_size for elem in it):
        raise RuntimeError('each element in list of batch should be of equal size')
    transposed = list(zip(*batch))  # It may be accessed twice, so we use a list.


    if isinstance(elem, tuple):
        return [collate(samples, collate_fn_map=collate_fn_map) for samples in transposed]  # Backwards compatibility.
    else:
        try:
            return elem_type([collate(samples, collate_fn_map=collate_fn_map) for samples in transposed])
        except TypeError:
            # The sequence type may not support `__init__(iterable)` (e.g., `range`).
            return [collate(samples, collate_fn_map=collate_fn_map) for samples in transposed]

https://github.com/pytorch/pytorch/blob/8349bf1cd1d5df7be73b194940bcf96209159f40/torch/utils/data/_utils/collate.py#L126-L149

I guess a proper solution for supporting empty batches would be to recycle this code but returning tuples/dicts/lists/... of empty tensors with torch.zeros, so that the types are still preserved, though filled with empty tensors.

Copy link
Contributor Author

@ffuuugor ffuuugor Oct 31, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That makes perfect sense, thanks! You're absolutely right this is the right way forward.

However, I don't see this as a blocker for landing this PR. This PR does solve the problem for the subset of cases and could be considered an atomic improvement. Not to delay merging it, I created an issue to track the proposed improvement (#534), hopefully someone will pick it up soon

if collate_fn is None:
collate_fn = default_collate
Expand All @@ -175,7 +176,11 @@ def __init__(
dataset=dataset,
batch_sampler=batch_sampler,
num_workers=num_workers,
collate_fn=wrap_collate_with_empty(collate_fn, sample_empty_shapes, dtypes),
collate_fn=wrap_collate_with_empty(
collate_fn=collate_fn,
sample_empty_shapes=sample_empty_shapes,
dtypes=dtypes,
),
pin_memory=pin_memory,
timeout=timeout,
worker_init_fn=worker_init_fn,
Expand Down