Can we switch to torchvision v2 transforms?

[cgebbe]

My understanding is that lightly uses the "old" v1 torchvision transforms, see e.g. https://github.com/lightly-ai/lightly/blob/8878b5b5e816bbd1f11b5c0a866eaee45a9a6004/lightly/transforms/dino_transform.py#L4

The v1 transforms have some drawbacks:

• only work on PIL images, not on torch tensors (which are the output of litdata.StreamingDataset)
• don't support segmentation datasets, see https://pytorch.org/blog/extending-torchvisions-transforms-to-object-detection-segmentation-and-video-tasks/#the-new-transforms-api

Would you accept a PR where we add the v2 versions? My understanding is that it's a drop-in replacement and that everything should stay the same, but I'd double check.

[IgorSusmelj]

This is something we already thought about. We should check if would require us to update the minimal dependencies. From https://github.com/lightly-ai/lightly/blob/master/requirements/minimal_requirements.txt

pytorch_lightning==1.7.1
torch==1.11.0
torchvision==0.12.0

I think torchvision introduced transforms v2 with 0.15.0 which would require us to bump both torch and torchvision.

[guarin]

Hi, adding support for v2 transforms sounds like a great idea. As @IgorSusmelj mentioned, lightly still supports older torchvision versions that do not yet support v2 transforms. I would suggest to add the following at the top of lightly.transforms.__init__:

try:
    # Use transforms v2 if available.
    import torchvision.transforms.v2 as torchvision_transforms
except ImportError:
    import torchvision.transforms as torchvision_transforms

And then replace code using the torchvision transforms in all files in lightly.transforms and lightly.collate:

# Replace this:
import torchvision.transforms as T
# With this:
import lightly.transforms.torchvision_transforms as T

This allows us to do the try/except check in a single place.

A PR would be very welcome (doesn't have to implement all the changes) :)