Avoid multiple pytree flattenings inside the prototype transforms

[pmeier]

By design, our transforms v2 can handle arbitrary input structures. Internally, we are using torch.utils._pytree for it.

Status quo

The simplest transforms flatten / unflatten only once:

vision/torchvision/prototype/transforms/_transform.py

Lines 37 to 42 in 54a2d4e

	flat_inputs, spec = tree_flatten(sample)
	flat_outputs = [
	self._transform(inpt, params) if _isinstance(inpt, self._transformed_types) else inpt
	for inpt in flat_inputs
	]
	return tree_unflatten(flat_outputs, spec)

However, albeit hidden, most transforms flatten at least twice if not multiple times:

• If a transform needs to know the spatial size to compute its params

  vision/torchvision/prototype/transforms/_geometry.py

  Line 108 in 54a2d4e

  height, width = query_spatial_size(sample)

  the extraction logic flattens again

  vision/torchvision/prototype/transforms/_utils.py

  Lines 101 to 102 in 54a2d4e

  	def query_spatial_size(sample: Any) -> Tuple[int, int]:
  	flat_sample, _ = tree_flatten(sample)

• If a transform performs some checks on the sample before transforming

  vision/torchvision/prototype/transforms/_geometry.py

  Lines 188 to 190 in 54a2d4e

  	if has_any(inputs, features.BoundingBox, features.Mask):
  	raise TypeError(f"BoundingBox'es and Mask's are not supported by {type(self).__name__}()")
  	return super().forward(*inputs)

  the checking utility flattens again

  vision/torchvision/prototype/transforms/_utils.py

  Lines 124 to 125 in 54a2d4e

  	def has_any(sample: Any, *types_or_checks: Union[Type, Callable[[Any], bool]]) -> bool:
  	flat_sample, _ = tree_flatten(sample)

This of course has some performance implications that can be avoided.

Proposal

In all of the cases where we perform the extra flattening are happening in internal and thus not user-facing methods. Thus, instead of keeping the option to operate on arbitrary input structures on our utilities, we could have them just work on already flattened inputs. This would avoid the repeated tree_flatten calls inside of them.

For our transforms that means two changes:

1. The extra calls in _get_params can be avoided by simply flattening before its call

   vision/torchvision/prototype/transforms/_transform.py

   Lines 35 to 37 in 54a2d4e

   	params = self._get_params(sample)
   	flat_inputs, spec = tree_flatten(sample)

2. The extra calls in overridden forward's require more boilerplate code. Basically, each transform that overrides forward needs to perform the flattening / unflattening themselves since the check utilities are called before the super().forward(...) call. This also means that we are technically still flattening twice, although the second time inside super().forward(...) does nothing.

   However, there is introduce _check method for type checks on prototype transforms #6503 that introduces a common interface for the checks. IIRC, we never followed up on it, since we eliminated some boilerplate in the overridden forward in [proto] Simplified code in overridden transform forward methods #6504. Since this proposal would re-add some boilerplate for performance gains, we could pick introduce _check method for type checks on prototype transforms #6503 up again. If we do, this leaves very few, objectively "outlier" transforms that would need to have this boiler plate.

   If we go for the common check interface, it could receive the already flattened sample as well similar to what was proposed in 1.

Conclusion

Flattening an input sample multiple times inside a single transformation has no benefits while slowing down execution. This issue proposes a way to avoid this while keeping the UI as convenient as it is.

cc @vfdev-5 @datumbox @bjuncek

[datumbox]

I'm agree with the proposal. Note that I have previously advocated against introducing more "protected" methods in the API, especially because it was actively developed. But given the simplicity of the refactoring and the fact it eliminates unnecessary calls, I'm OK to proceed with this. IMO anything that will eliminate unnecessary CPU cycles and doesn't change the public interface, is worth considering.

[pmeier]

I've run some benchmarks:

from time import perf_counter_ns

import torch
from torch.utils._pytree import tree_flatten
from torchvision.prototype import features


def timeit(
    descr, sample, *, num_transforms, num_calls_per_sample, num_samples, n=10_000
):
    time_diffs = []
    for _ in range(n):
        start = perf_counter_ns()
        tree_flatten(sample)
        stop = perf_counter_ns()
        time_diffs.append(stop - start)

    overhead_per_call = (
        torch.tensor(time_diffs, dtype=torch.float64).mul_(1e-9).median()
    )
    estimated_overhead_per_epoch = (
        overhead_per_call * num_calls_per_sample * num_samples
    )
    estimated_overhead_per_epoch_with_proposal = (
        overhead_per_call * num_transforms * num_samples
    )

    print(f"{descr}\n")
    print(f"Overhead per call: {overhead_per_call * 1e6:.1f} µs")
    print(f"Estimated overhead per epoch: {estimated_overhead_per_epoch / 60:.1f} m")
    print(
        f"Estimated overhead per epoch with proposal: {estimated_overhead_per_epoch_with_proposal / 60:.1f} m"
    )
    print("-" * 80 + "\n")


timeit(
    "Image classification",
    (features.Image(torch.rand(3, 512, 512)), features.Label(0)),
    # This assumes all possible transforms in the references are used
    num_transforms=9,
    num_calls_per_sample=18,
    # num samples per imagenet epoch * number of epochs / number of processes
    num_samples=1_200_000 * 600 / 8,
)

timeit(
    "Object detection",
    (
        features.Image(torch.rand(3, 512, 512)),
        dict(
            boxes=features.BoundingBox(
                torch.rand(8, 4),
                format=features.BoundingBoxFormat.XYXY,
                spatial_size=(512, 512),
            ),
            labels=features.Label(torch.randint(0, 10, (8,))),
            masks=features.Mask(torch.randint(0, 256, (512, 512), dtype=torch.uint8)),
            image_id=torch.tensor(0),
            area=torch.rand(8),
            iscrowd=torch.randint(0, 2, (8,)),
        ),
    ),
    # This assumes lsj augmentation with copy-paste
    num_transforms=5,
    num_calls_per_sample=10,
    # num samples per coco epoch * number of epochs / number of processes
    num_samples=80_000 * 600 / 8,
)

Overhead per call: 3.0 µs
Estimated overhead per epoch: 82.2 m
Estimated overhead per epoch with proposal: 41.1 m
--------------------------------------------------------------------------------

Object detection

Overhead per call: 9.2 µs
Estimated overhead per epoch: 9.2 m
Estimated overhead per epoch with proposal: 4.6 m
--------------------------------------------------------------------------------

• If the estimation is correct, adopting this can shave roughly 40 minutes from the classification training
• The reduction for objection detection is less significant. However, due to the more complex sample structure, flattening the sample once is roughly 3x slower than an classification sample. Since the datasets v2 will in general return more rather than less information, this proposal will help us in the future as well.

[pmeier]

We can achieve even less overhead if we find a protocol for a container transform to properly communicate to its children whether the input is already flattened or not. For example, a transforms.Compose could flatten initially, let all children do their magic on the flattened sample without trying to flatten again, and only unflatten after all children had their turn.

This can work if the sample structure doesn't change over the course of the pipeline. I think so far this is the case for our transforms.

[pmeier]

I think the most convenient would be to communicate this through the type:

class FlatInputs(list):
    pass

# inside Compose.forward()

flat_inputs, spec = tree_flatten(inputs)
for transform in self.transforms:
    transform(FlatInputs(flat_inputs))
return tree_unflatten(flat_inputs, spec)

# inside a child Transform.forward()

flat_inputs, spec = (inputs, None) if isinstance(inputs, FlatInputs) else tree_flatten(inputs)

flat_outputs = ...

return flat_outputs if spec is None else tree_unflatten(flat_outputs, spec)

[vfdev-5]

Maybe, we can also see FlatInputs as a placeholder for flat_inputs and spec as attributes. This way FlatInputs wont be just a dummy class with a name to distinguish.

[pmeier]

You mean like

class FlatInputs:
    def __init__(self, flat_inputs):
        self.flat_inputs = flat_inputs

    def __iter__(self):
        yield self.flat_inputs
        yield LeafSpec()

# in Transform.forward

flat_inputs, spec = inputs if isinstance(inputs, FlatInputs) else tree_flatten(inputs)

flat_outputs = ...

return tree_unflatten(flat_outputs, spec)

?

[vfdev-5]

class FlatInputs:
    def __init__(self, flat_inputs, spec):
        self.flat = flat_inputs
        self.spec = spec


# inside a child Transform.forward()

flat_inputs, spec = (inputs.flat, inputs.spec) if isinstance(inputs, FlatInputs) else tree_flatten(inputs)

flat_outputs = ...

return flat_outputs if isinstance(inputs, FlatInputs) else tree_unflatten(flat_outputs, spec)

[pmeier]

I played with it, but my as well as your proposal has one major flaw: it is not a list. That sounds pedantic, but I don't think it is. With class FlatInputs(list): the input to forward is always a list. Thus, users that overwrite the forward in their custom transformation don't have to learn yet another idiom how to parse the input into something they can work with.

[pmeier]

Found an issue with the Compose proposal: we can only do this if every single child is able to handle flattened inputs. This works for all our builtin transforms, but what if someone uses children like

vision/test/test_prototype_transforms.py

Line 1199 in dc5fd83

[lambda x: 2.0 * x, transforms.Pad(2), transforms.RandomCrop(28)],

Passing a flattened input to the lambda will fail.

In general, I don't think there is a way of knowing whether or not the children transforms can handle flattened inputs or not. I currently only see two options here:

1. Add a __handles_flat_inputs__ = True attribute to every transformation that does, i.e. all of our builtin transformations. With that in place, a container transform could check the children and if all come back positive, we use the approach detailed above. Otherwise, we let each transform handle the input themself. If we want users to also be able to mark their custom transforms, we should probably make this a public attribute.
2. Add a flatten_once: bool = False to the container transforms so the user needs to opt in if they want the speed-up.

[pmeier]

As discussed offline, we are going to move forward with the original proposal and open another issue for the container proposal since that is more controversial.

Edit: See #6769.