[FEEDBACK] Transforms V2 API

[datumbox]

🚀 The feature

This issue is dedicated for collecting community feedback on the Transforms V2 API. Please review the dedicated blogpost where we describe the API in detail and provide an overview of its features.

We would love to get your thoughts, comments and input in order to improve the API and graduate it from prototype on the near future.

Please also check out #7319 where we collect feedback on some specific design decision, and document as well which APIs may change in the future!

---

Code example using this image:

import PIL
from torchvision import io, utils
from torchvision.prototype import features, transforms as T
from torchvision.prototype.transforms import functional as F


# Defining and wrapping input to appropriate Tensor Subclasses
path = "COCO_val2014_000000418825.jpg"
img = features.Image(io.read_image(path), color_space=features.ColorSpace.RGB)
# img = PIL.Image.open(path)
bboxes = features.BoundingBox(
    [[2, 0, 206, 253], [396, 92, 479, 241], [328, 253, 417, 332],
     [148, 68, 256, 182], [93, 158, 170, 260], [432, 0, 438, 26],
     [422, 0, 480, 25], [419, 39, 424, 52], [448, 37, 456, 62],
     [435, 43, 437, 50], [461, 36, 469, 63], [461, 75, 469, 94],
     [469, 36, 480, 64], [440, 37, 446, 56], [398, 233, 480, 304],
     [452, 39, 463, 63], [424, 38, 429, 50]],
    format=features.BoundingBoxFormat.XYXY,
    spatial_size=F.get_spatial_size(img),
)
labels = features.Label([59, 58, 50, 64, 76, 74, 74, 74, 74, 74, 74, 74, 74, 74, 50, 74, 74])


# Defining and applying Transforms V2
trans = T.Compose(
    [
        T.ColorJitter(contrast=0.5),
        T.RandomRotation(30),
        T.CenterCrop(480),
    ]
)
img, bboxes, labels = trans(img, bboxes, labels)


# Visualizing results
viz = utils.draw_bounding_boxes(F.to_image_tensor(img), boxes=bboxes)
F.to_pil_image(viz).show()

[rsokl]

Could you post the link to the blog post in this thread when it becomes available? I am not sure where else to look to find it.

[jangop]

Can you provide a good starting point for getting an overview of the current state of Transforms V2?

Or is reading through https://github.com/pytorch/vision/tree/main/torchvision/prototype/transforms and https://github.com/pytorch/vision/projects/5 the best approach for now?

[datumbox]

@rsokl @jangop We got a blogpost in the pipeline that provides an overview. We are waiting for marketing to publish. I'll post the link here once we do. Until then those 2 references are the best places to look. There is no documentation because the API was being modified but all transforms receive the same exact parameters as V1.

[datumbox]

@rsokl @jangop we published the blogpost. Looking forward to your input.

[rsokl]

Hi all,

I just read the blog post and have been exploring the prototype's internals for the last few of days: torchvision.transforms v2 looks great! I'd like to share my thoughts with you all.

First, let me summarize some of the primary features that are offered by v2. This will, in part, help me make sure that I have a clear understanding of things. v2 includes:

• Native support for transforming videos, bounding boxes, segmentation masks, and labels via the _Feature class.
• The ability to describe metadata for these features. E.g., the image class that has a mechanism for signaling the image's color space (e.g. RGB vs RGBA).
• A flexible API for mapping transformations over general pytrees of tensors / features, making it easy to incorporate arbitrary numbers/configurations of features in each transformation.
• A simple API for implementing transforms via the Transform class. Here, an implementation draws random parameters within _get_params, and an individual tensor/feature is transformed by _transform.
• Continued support for batches of tensors/features as well as cuda support.
• Improved performance over v1.

I had been doing a review of augmentation libraries, including kornia, albumentations, and augly, when I came across v2 for torchvision. To me, your new API is the simplest, most capable, and the easiest to extend. I particularly like how simple Transform is and that you do not rely on sprawling class hierarchies to dispatch functionality.

So first of all, thank you for all of the hard work that you have been doing on this effort. The PyTorch community is fortunate to benefit from this excellent work ❤️.

There are two features that I would like to propose. I would be happy to assist with these if there is interest.

Enabling local reproducibility by passing torch.Generator to transforms.

In recent years, NumPy has completely revised their PRNG API to avoid global random state (here is a great post on good practices with NumPy's generators). JAX avoids mutable RNG objects altogether. PyTorch provides torch.Generator to users to to make randomness local and "non-spooky", but many libraries prevent users from utilizing this capability.

I am proposing that Transform enable users to optionally pass in a Generator to the forward pass so that torchvision transform pipelines can be made to be isolated from global entropy and thus support more reproducible workflows. This reproducibility is especially useful in the context of performing testing & evaluation – the specific sequence of data transformations performed should be able to be isolated from whether or not a model is using dropout in its forward pass.

The local generator (whose default value is PyTorch's global generator) could be passed to the forward method and piped to the _get_params method. This assumes that all random number generation occurs in self._get_params. The following would be a compatibility-preserving modification to Transform; by default the global seed would still control stochasticity in the transforms.

from torch import Generator, Tensor, default_generator
import torch.nn as nn

class Transform(nn.Module):
    def _get_params(self, flat_inputs: List[Any], *, generator: Generator = default_generator) -> Dict[str, Any]:
        ...

    def forward(self, *inputs: Any, generator: Generator = default_generator) -> Any:
        # the only modification
        params = self._get_params(flat_inputs, generator=generator)

and transforms that implement _get_params would replace calls like

# e.g. replace calls like
angle = float(torch.empty(1).uniform_(0.0, 180.).item())

with

# specifying the device is, unfortunately, necessary: https://github.com/pytorch/pytorch/issues/79018
angle = float(torch.empty(1, device=generator.device).uniform_(0.0, 180., generator=generator).item())

A transform like Compose would have to be modified as well. Currently, it supports a sequence of callables that are assumed to accept a single positional argument. It could be assumed that only instances of Transform involve stochasticity and will be passed the random generator. In this case, Compose would look like:

class Compose(Transform):
    # __init__ is unchanged

    def forward(self, *inputs: Any, generator: Generator = default_generator) -> Any:
        sample = inputs if len(inputs) > 1 else inputs[0]
        for transform in self.transforms:
            sample = transform(sample) if not isinstance(transform, Transform) else transform(sample, generator=generator)
        return sample

It would be straightforward to document this behavior to users – that only instances of Transform are passed the generator – so that they know how to opt-in to having the generator be passed to their custom transforms. And, again, this would be compatible with the old nn.Module transforms.

An example of this in practice would be:

from torch import Generator

rng = Generator.manual_seed(0)

trans = T.Compose(
    [
        T.ColorJitter(contrast=0.5),
        T.RandomRotation(30),
        T.CenterCrop(480),
    ]
)
img, bboxes, labels = trans(img, bboxes, labels, generator=rng)

Another nice thing about this is that specific fail cases that occur during training/testing can be reproduced in an isolated way; _get_params(dummy_img, generator=rng) can be used to iterate the generator's state to "replay" a sequence of transformations without have to redo all of the compute. Whereas this would not work if the model and the transforms both affect and derive from global state.

Adding Feature to the public API and adding a protocol for dispatching transforms on 3rd party features

This request a bit more ambitious, but I think that it could have a big pay-off.

The crown jewel of torchvision.transforms v2 is its added support for features like bounding boxes and segmentation masks. Presently, the development of new features and transforms is gated on the development efforts of the intrepid torchvision team. It would be great to enable users to implement their own features (e.g. Polygon or ImageTrack) and for torchvision's transforms (or 3rd party transforms) to be able to dispatch on those features. This would reduce the burden on the torchvision team in that it enables users to develop bespoke features that have "drop-in" compatibility with standard torchvision pipelines, without you needing to officially support the feature.

Thus I am proposing to a) make _Feature part of the public API (take away that leading underscore 😄 ) so that 3rd parties can create custom features, and b) provide a dispatch mechanism by which a 3rd party feature can coerce a transform (either 1st party or 3rd party transform) into dispatching to its own implementation of said transform. E.g. I could write a Polygon feature that would be able to implement its own affine transformation that both torchvision.transforms.affine and torchvision.transforms.RandomAffine dispatch to when they operate on a Polygon tensor.

There is already some degree of dispatching going on among some transforms. E.g., for affine and RandomAffine, any _Feature subclass can simply implement a Feature.affine method. The affine function checks isinstance(inp, Feature) and then calls inp.affine(...). And RandomAffine._transform simply calls affine under the hood, which performs the aforementioned dispatching.

That being said, users can't rely on all functions and transformations to be implemented this way. Indeed, Transform._transformed_types will often preclude such synergistic dispatching (RandomAffine._transformed_types works here because it only contains Tensor). It seems to me that, to support this proposal, Transform._transformed_types would need to be replaced by a Feature-centric way of documenting compatibility with transforms.

Given that these torchvision's transforms are all unary functions, I believe that the dispatch mechanism could be kept relatively simple. That being said, I don't have deep insight into how JIT compatibility might impact the feasibility of some of the design ideas that I have been playing around with. I also only have a rudimentary understanding of __torch_function__ and __torch_dispatch__, but I get the sense that torchvision's transforms would not be able to leverage these.

Ultimately, I do not have a concrete design for this dispatch protocol, but would be happy to help develop it with you.

Thank you very much for your consideration. And thank you again for your excellent work on v2. My colleagues and I are excited to start using it.

[pmeier]

Thanks a ton for the kind words @rsokl! I think your summary is on point. Let me share my thoughts on your proposals.

Allow torch.Generator

+1 from me on basically everything you proposed. I think the design is simple enough that this can be done quickly in case we decide to adopt it. pytorch/pytorch#79018 is unfortunate here, but if we as torchvision need that, we might be able to speed things up on core.

The only thing that I would change about your proposal is the existence of a default value for generator on _get_params. During normal operation, _get_params should not be called directly and thus we can be explicit about it. IIUC, this would also not interfere with your use case of fast-forwarding the RNG without doing the actual computation since there you also don't want to use the default generator.

Regarding this use case, please be aware that it will only work for most transforms. Some more elaborate transforms like the AutoAugment family or SimpleCopyPaste don't fit into the regular Transform scheme that we devised. They overwrite forward directly and perform the random parameter generation there.

Make features._Feature public

Let's go through the proposal point by point

a) make _Feature part of the public API (take away that leading underscore 😄 ) so that 3rd parties can create custom features

Opinions about this changed over the course of the development. We started with it being public, but made it private at some point. Recently, we found a use case for it to be public again (#6663 (comment)). A cautious +1 from my side on this.

b) provide a dispatch mechanism by which a 3rd party feature can coerce a transform (either 1st party or 3rd party transform) into dispatching to its own implementation of said transform.

As you have explained above, this is already happening for some dispatchers like F.affine. However, there are a few caveats that need more careful design:

1. Not all dispatchers are available as methods on the features. In the beginning our rule of thumb has been we only put something as method if in addition to images, bounding boxes and masks are supported as well. This is not a design limitation though, so we only need to "fix" our rule if we want this.

2. Not all kernels have dispatchers. For example

   vision/torchvision/prototype/transforms/functional/_meta.py

   Line 161 in 4f3a000

   def convert_format_bounding_box(

   because there would only be a single kernel. Even if we added such a dispatcher for consistency we would get into hot waters, since the dispatcher would violate the rule "for BC plain tensors need to be treated as image or video" since there is no image or video kernel to dispatch to. This rule is currently kept by all our other dispatchers which allows us to keep them BC even with regard to JIT. Of course there is no BC concern for new functions, but it is a lot easier to say "the functional API is JIT compatible" rather than "the functional API that was already there for v1 is JIT compatible, but for the remainder it depends".

3. Internally, a lot of transformations look up the spatial size from the inputs with

   vision/torchvision/prototype/transforms/_geometry.py

   Line 353 in 4f3a000

   height, width = query_spatial_size(flat_inputs)

   This is already pretty flexible and can pull the information from a lot of types, but there is no "official" protocol yet. For now we are using the spatial_size attribute

   vision/torchvision/prototype/transforms/functional/_meta.py

   Lines 103 to 104 in 4f3a000

   	elif isinstance(inpt, (features.Image, features.Video, features.BoundingBox, features.Mask)):
   	return list(inpt.spatial_size)

   but this only a convention.

4. Taking 3. one step further, some transformations also query bounding boxes like that:

   vision/torchvision/prototype/transforms/_utils.py

   Lines 105 to 108 in 4f3a000

   	def query_bounding_box(flat_inputs: List[Any]) -> features.BoundingBox:
   	bounding_boxes = [inpt for inpt in flat_inputs if isinstance(inpt, features.BoundingBox)]
   	if not bounding_boxes:
   	raise TypeError("No bounding box was found in the sample")

   Without a more elaborate protocol, any bounding-box-like feature like RotatedBoundingBox proposed in [RFC] Rotated Bounding Boxes #2761 cannot be a free feature, but has to subclass features.BoundingBox. Otherwise it won't be picked up and in turn the transformation will fail.

These are just the limitations from the top of my head. There will likely be more if take a closer look. Thus, the design for this will be more complex.

That being said some of these points mentioned above are already in our backlog, but we focused on performance the past few weeks. I agree allowing users to extend our API like that would be very powerful, but we need to solve all rough edges before we can promote this as official part of the API. Let's hear what @datumbox and @vfdev-5 say about this before I cast my vote.

[jangop]

I had been doing a review of augmentation libraries, including kornia, albumentations, and augly, when I came across v2 for torchvision. To me, your new API is the simplest, most capable, and the easiest to extend. I particularly like how simple Transform is and that you do not rely on sprawling class hierarchies to dispatch functionality.

@rsokl Would you perhaps share a few words on why / how / in which regard Transforms V2 is “the simplest, most capable, and the easiest to extend”? Reading through the blog post, the one section I missed was related work for context. Personally, I only have experience with albumentations. Coming from that, I am inclined to agree with your assessment.

[FlorinAndrei]

Is there an estimate in terms of either time, or version number, or both, for when the v2 transforms will be included in pytorch-stable? The blog says "planned for Q1" but I wonder if there's a better estimate than that somewhere.

I was about to implement, on a small scale, functions that basically do what v2 does, but if your prototype works well then I will use it for my project.

Should I expect any issues if I train a model on the nightly builds with the v2 transforms, but then run it for predictions on pytorch-stable 1.13? Other than the v2 transforms, I do not use any unusual features.

[FlorinAndrei]

test case: image segmentation with SegFormer. My dataset has images and masks organized in tuples. I am following this example:

https://huggingface.co/blog/fine-tune-segformer

I need to make sure that geometric transforms (horizontal flip, rotation, affine, perspective) are applied randomly, but are applied the exact same way to each image and its corresponding mask. I am testing with horizontal flip and the image and the mask are flipped in an uncoordinated fashion.

Code example (admittedly inefficient):

from torchvision.prototype.transforms import (
    Compose,
    RandomApply,
    ColorJitter,
    RandomRotation,
    RandomCrop,
    RandomAffine,
    RandomHorizontalFlip,
    RandomPerspective,
)
from torchvision.prototype import features
from torchvision.prototype.features import ColorSpace

feature_extractor = SegformerFeatureExtractor()

augmentations = Compose(
    [
        # RandomApply([ColorJitter(brightness=0.5, contrast=0.5)], p=0.75),
        RandomHorizontalFlip(p=0.5),
        # RandomApply([RandomRotation(degrees=45)], p=0.75),
        # RandomApply([RandomAffine(degrees=0.0, scale=(0.5, 2.0))], p=0.25),
        # RandomPerspective(distortion_scale=0.5, p=0.25),
    ]
)


def train_transforms(example_batch):
    # original
    images = [augmentations(x) for x in example_batch["pixel_values"]]
    labels = [x for x in example_batch["label"]]
    inputs = feature_extractor(images, labels)
    return inputs


def train_transforms2(example_batch):
    # my version
    batch_items = list(
        zip(
            [x for x in example_batch["pixel_values"]],
            [x for x in example_batch["label"]],
        )
    )
    batch_items_aug = [
        augmentations(
            features.Image(
                np.swapaxes(np.array(x[0]), 0, -1), color_space=ColorSpace.RGB
            ),
            features.Mask(np.swapaxes(np.array(x[1]), 0, -1)),
        )
        for x in batch_items
    ]
    images, labels = map(list, zip(*batch_items_aug))
    inputs = feature_extractor(images, labels)
    return inputs


train_ds.set_transform(train_transforms2)

This is one entry in the train dataset that the transforms are applied to:

{'dataset': 0,
 'pixel_values': <PIL.PngImagePlugin.PngImageFile image mode=RGB size=400x322>,
 'tumor': 1,
 'dataset_tumor': 1,
 'label': <PIL.PngImagePlugin.PngImageFile image mode=L size=400x322>}

If I repeatedly visualize the image and the mask in train_ds[0], they are flipped randomly, but the flip is not applied the same to the image and its mask - they are each flipped as a completely separate random process.

What I expect:

• images and masks are geometrically transformed in a coordinated fashion (random transforms are applied the same way to an image and to its mask)
• pixel-value transforms such as brightness are only applied to the image (this seems to work fine)

[pmeier]

Hey @FlorinAndrei

Is there an estimate in terms of either time, or version number, or both, for when the v2 transforms will be included in pytorch-stable? The blog says "planned for Q1" but I wonder if there's a better estimate than that somewhere.

We are aiming to publish it with the next release, i.e. torchvision==0.15.0. The exact date is not yet available. We operate on a 3-4 months release cycle and the last release just happened on the 28th of October. So you can expect this to be released in February or March 2023.

Should I expect any issues if I train a model on the nightly builds with the v2 transforms, but then run it for predictions on pytorch-stable 1.13? Other than the v2 transforms, I do not use any unusual features.

No. Apart from JIT, the transforms are BC and thus there should be no complications. You can have a look at #6433 where we do just that.

I am testing with horizontal flip and the image and the mask are flipped in an uncoordinated fashion.

You are not flipping the masks at all:

images = [augmentations(x) for x in example_batch["pixel_values"]]
labels = [x for x in example_batch["label"]]

The horizontal flip is happening in augmentations but this is never called on any example_batch["label"]. Even if you were using it in the line above, the augmentation cannot be same for images and masks since you apply them separately.

Transforms v2 can handle arbitrary input structures and so you don't need to handle images and masks separately. You can just pass them into the transforms together like:

augmented_batch = [augmentations(sample) for sample in example_batch]
images, labels = zip(*augmented_batch)
inputs = feature_extractor(images, labels)

The only caveat of this is that example_batch["label"] will be treated as image for BC. Meaning, the color transformations will also be applied to it. Thus, we need to communicate to the transformations that it is a mask.

In transforms v2 this is done by wrapping the data into custom tensor classes located under torchvision.prototype.features:

from torchvision.prototype import features, transforms
from torchvision.prototype.transforms import functional as F

class WrapIntoFeatures(transforms.Transform):
    def forward(self, sample):
        image, label = sample
        label = features.Mask(F.pil_to_tensor(mask))
        return image, label

augmentations = transforms.Compose([WrapIntoFeatures(), ...])

For custom datasets this will have to happen manually. For our builtin ones, we currently explore the options. See #6662 for the current favorite.

I've put together a small notebook for you. I hope it helps. but let us know if you encounter anything.

[maxoppelt]

I understand the need of wrapping the torch.Tensor class to mimic the behavior of "matching" transformations for a specific input type using _transformed_types class variable. However, as far as I can understand the user directly handles / sets these types. This might lead to inconsistent behavior with functionalities commonly used and provided by inheritance of the th.Tensor class. I have attached a simple multiplication of two image-typed-like-feature classes below to highlight my concerns. I see two directions:

(a) It is fine for "transforms v2", that operation on two feature.Image types as input defaults back to a non features.Image tensor. I think this behavior might feel inconsistent for some users?
(b) The operations (multiplication, addition, substraction, etc.) need to be implemented. However, this opens the hells gate to a lot of manual implementation overhead, e.g. what is the multiplication of a black-and-white image and a RGB image. Or what is the multiplication of an features.Image and a features.Video? --> Implementation Hell. Not to mention what would happen if a feature.Image would be passed to a neural network... feature.Image * th.Tensor..

It might be that my concerns are irrelevant: I have just read through the code briefly. However, I think that this behavior needs to be specified / discussed upfront.

import torch as th
from torchvision.prototype.features import Image

img1 = Image(th.rand(3, 256, 256))
img2 = Image(th.rand(3, 256, 256))

type(img1), type(img2)
>> torchvision.prototype.features._image.Image, torchvision.prototype.features._image.Image

img = img1 * img2
type(img)
>> torch.Tensor