Skip to content

Commit

Permalink
tv_tensor -> TVTensor where it matters (pytorch#7904)
Browse files Browse the repository at this point in the history
Co-authored-by: Philip Meier <[email protected]>
  • Loading branch information
NicolasHug and pmeier authored Aug 30, 2023
1 parent d5f4cc3 commit 25ec3f2
Show file tree
Hide file tree
Showing 9 changed files with 84 additions and 65 deletions.
2 changes: 1 addition & 1 deletion docs/source/transforms.rst
Original file line number Diff line number Diff line change
Expand Up @@ -183,7 +183,7 @@ Transforms are available as classes like
This is very much like the :mod:`torch.nn` package which defines both classes
and functional equivalents in :mod:`torch.nn.functional`.

The functionals support PIL images, pure tensors, or :ref:`tv_tensors
The functionals support PIL images, pure tensors, or :ref:`TVTensors
<tv_tensors>`, e.g. both ``resize(image_tensor)`` and ``resize(bboxes)`` are
valid.

Expand Down
9 changes: 5 additions & 4 deletions docs/source/tv_tensors.rst
Original file line number Diff line number Diff line change
Expand Up @@ -5,10 +5,11 @@ TVTensors

.. currentmodule:: torchvision.tv_tensors

TVTensors are tensor subclasses which the :mod:`~torchvision.transforms.v2` v2 transforms use under the hood to
dispatch their inputs to the appropriate lower-level kernels. Most users do not
need to manipulate tv_tensors directly and can simply rely on dataset wrapping -
see e.g. :ref:`sphx_glr_auto_examples_transforms_plot_transforms_e2e.py`.
TVTensors are :class:`torch.Tensor` subclasses which the v2 :ref:`transforms
<transforms>` use under the hood to dispatch their inputs to the appropriate
lower-level kernels. Most users do not need to manipulate TVTensors directly and
can simply rely on dataset wrapping - see e.g.
:ref:`sphx_glr_auto_examples_transforms_plot_transforms_e2e.py`.

.. autosummary::
:toctree: generated/
Expand Down
4 changes: 2 additions & 2 deletions gallery/transforms/plot_custom_transforms.py
Original file line number Diff line number Diff line change
Expand Up @@ -74,7 +74,7 @@ def forward(self, img, bboxes, label): # we assume inputs are always structured
print(f"Output image shape: {out_img.shape}\nout_bboxes = {out_bboxes}\n{out_label = }")
# %%
# .. note::
# While working with tv_tensor classes in your code, make sure to
# While working with TVTensor classes in your code, make sure to
# familiarize yourself with this section:
# :ref:`tv_tensor_unwrapping_behaviour`
#
Expand Down Expand Up @@ -111,7 +111,7 @@ def forward(self, img, bboxes, label): # we assume inputs are always structured
# In brief, the core logic is to unpack the input into a flat list using `pytree
# <https://github.com/pytorch/pytorch/blob/main/torch/utils/_pytree.py>`_, and
# then transform only the entries that can be transformed (the decision is made
# based on the **class** of the entries, as all tv_tensors are
# based on the **class** of the entries, as all TVTensors are
# tensor-subclasses) plus some custom logic that is out of score here - check the
# code for details. The (potentially transformed) entries are then repacked and
# returned, in the same structure as the input.
Expand Down
6 changes: 3 additions & 3 deletions gallery/transforms/plot_custom_tv_tensors.py
Original file line number Diff line number Diff line change
@@ -1,14 +1,14 @@
"""
=====================================
====================================
How to write your own TVTensor class
=====================================
====================================
.. note::
Try on `collab <https://colab.research.google.com/github/pytorch/vision/blob/gh-pages/main/_generated_ipynb_notebooks/plot_custom_tv_tensors.ipynb>`_
or :ref:`go to the end <sphx_glr_download_auto_examples_transforms_plot_custom_tv_tensors.py>` to download the full example code.
This guide is intended for advanced users and downstream library maintainers. We explain how to
write your own tv_tensor class, and how to make it compatible with the built-in
write your own TVTensor class, and how to make it compatible with the built-in
Torchvision v2 transforms. Before continuing, make sure you have read
:ref:`sphx_glr_auto_examples_transforms_plot_tv_tensors.py`.
"""
Expand Down
14 changes: 7 additions & 7 deletions gallery/transforms/plot_transforms_getting_started.py
Original file line number Diff line number Diff line change
Expand Up @@ -115,7 +115,7 @@
# segmentation, or videos (:class:`torchvision.tv_tensors.Video`), we could have
# passed them to the transforms in exactly the same way.
#
# By now you likely have a few questions: what are these tv_tensors, how do we
# By now you likely have a few questions: what are these TVTensors, how do we
# use them, and what is the expected input/output of those transforms? We'll
# answer these in the next sections.

Expand All @@ -126,15 +126,15 @@
# What are TVTensors?
# --------------------
#
# TVTensors are :class:`torch.Tensor` subclasses. The available tv_tensors are
# TVTensors are :class:`torch.Tensor` subclasses. The available TVTensors are
# :class:`~torchvision.tv_tensors.Image`,
# :class:`~torchvision.tv_tensors.BoundingBoxes`,
# :class:`~torchvision.tv_tensors.Mask`, and
# :class:`~torchvision.tv_tensors.Video`.
#
# TVTensors look and feel just like regular tensors - they **are** tensors.
# Everything that is supported on a plain :class:`torch.Tensor` like ``.sum()``
# or any ``torch.*`` operator will also work on a tv_tensor:
# or any ``torch.*`` operator will also work on a TVTensor:

img_dp = tv_tensors.Image(torch.randint(0, 256, (3, 256, 256), dtype=torch.uint8))

Expand All @@ -146,7 +146,7 @@
# transform a given input, the transforms first look at the **class** of the
# object, and dispatch to the appropriate implementation accordingly.
#
# You don't need to know much more about tv_tensors at this point, but advanced
# You don't need to know much more about TVTensors at this point, but advanced
# users who want to learn more can refer to
# :ref:`sphx_glr_auto_examples_transforms_plot_tv_tensors.py`.
#
Expand Down Expand Up @@ -234,9 +234,9 @@
# Torchvision also supports datasets for object detection or segmentation like
# :class:`torchvision.datasets.CocoDetection`. Those datasets predate
# the existence of the :mod:`torchvision.transforms.v2` module and of the
# tv_tensors, so they don't return tv_tensors out of the box.
# TVTensors, so they don't return TVTensors out of the box.
#
# An easy way to force those datasets to return tv_tensors and to make them
# An easy way to force those datasets to return TVTensors and to make them
# compatible with v2 transforms is to use the
# :func:`torchvision.datasets.wrap_dataset_for_transforms_v2` function:
#
Expand All @@ -246,7 +246,7 @@
#
# dataset = CocoDetection(..., transforms=my_transforms)
# dataset = wrap_dataset_for_transforms_v2(dataset)
# # Now the dataset returns tv_tensors!
# # Now the dataset returns TVTensors!
#
# Using your own datasets
# ^^^^^^^^^^^^^^^^^^^^^^^
Expand Down
42 changes: 21 additions & 21 deletions gallery/transforms/plot_tv_tensors.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,18 +9,18 @@
TVTensors are Tensor subclasses introduced together with
``torchvision.transforms.v2``. This example showcases what these tv_tensors are
``torchvision.transforms.v2``. This example showcases what these TVTensors are
and how they behave.
.. warning::
**Intended Audience** Unless you're writing your own transforms or your own tv_tensors, you
**Intended Audience** Unless you're writing your own transforms or your own TVTensors, you
probably do not need to read this guide. This is a fairly low-level topic
that most users will not need to worry about: you do not need to understand
the internals of tv_tensors to efficiently rely on
the internals of TVTensors to efficiently rely on
``torchvision.transforms.v2``. It may however be useful for advanced users
trying to implement their own datasets, transforms, or work directly with
the tv_tensors.
the TVTensors.
"""

# %%
Expand All @@ -31,8 +31,8 @@


# %%
# What are tv_tensors?
# --------------------
# What are TVTensors?
# -------------------
#
# TVTensors are zero-copy tensor subclasses:

Expand All @@ -46,31 +46,31 @@
# Under the hood, they are needed in :mod:`torchvision.transforms.v2` to correctly dispatch to the appropriate function
# for the input data.
#
# :mod:`torchvision.tv_tensors` supports four types of tv_tensors:
# :mod:`torchvision.tv_tensors` supports four types of TVTensors:
#
# * :class:`~torchvision.tv_tensors.Image`
# * :class:`~torchvision.tv_tensors.Video`
# * :class:`~torchvision.tv_tensors.BoundingBoxes`
# * :class:`~torchvision.tv_tensors.Mask`
#
# What can I do with a tv_tensor?
# -------------------------------
# What can I do with a TVTensor?
# ------------------------------
#
# TVTensors look and feel just like regular tensors - they **are** tensors.
# Everything that is supported on a plain :class:`torch.Tensor` like ``.sum()`` or
# any ``torch.*`` operator will also work on tv_tensors. See
# any ``torch.*`` operator will also work on TVTensors. See
# :ref:`tv_tensor_unwrapping_behaviour` for a few gotchas.

# %%
# .. _tv_tensor_creation:
#
# How do I construct a tv_tensor?
# -------------------------------
# How do I construct a TVTensor?
# ------------------------------
#
# Using the constructor
# ^^^^^^^^^^^^^^^^^^^^^
#
# Each tv_tensor class takes any tensor-like data that can be turned into a :class:`~torch.Tensor`
# Each TVTensor class takes any tensor-like data that can be turned into a :class:`~torch.Tensor`

image = tv_tensors.Image([[[[0, 1], [1, 0]]]])
print(image)
Expand All @@ -92,7 +92,7 @@
print(image.shape, image.dtype)

# %%
# Some tv_tensors require additional metadata to be passed in ordered to be constructed. For example,
# Some TVTensors require additional metadata to be passed in ordered to be constructed. For example,
# :class:`~torchvision.tv_tensors.BoundingBoxes` requires the coordinate format as well as the size of the
# corresponding image (``canvas_size``) alongside the actual values. These
# metadata are required to properly transform the bounding boxes.
Expand All @@ -109,7 +109,7 @@
# ^^^^^^^^^^^^^^^^^^^^^^^^^^^
#
# You can also use the :func:`~torchvision.tv_tensors.wrap` function to wrap a tensor object
# into a tv_tensor. This is useful when you already have an object of the
# into a TVTensor. This is useful when you already have an object of the
# desired type, which typically happens when writing transforms: you just want
# to wrap the output like the input.

Expand All @@ -125,7 +125,7 @@
# .. _tv_tensor_unwrapping_behaviour:
#
# I had a TVTensor but now I have a Tensor. Help!
# ------------------------------------------------
# -----------------------------------------------
#
# By default, operations on :class:`~torchvision.tv_tensors.TVTensor` objects
# will return a pure Tensor:
Expand All @@ -151,7 +151,7 @@
# But I want a TVTensor back!
# ^^^^^^^^^^^^^^^^^^^^^^^^^^^
#
# You can re-wrap a pure tensor into a tv_tensor by just calling the tv_tensor
# You can re-wrap a pure tensor into a TVTensor by just calling the TVTensor
# constructor, or by using the :func:`~torchvision.tv_tensors.wrap` function
# (see more details above in :ref:`tv_tensor_creation`):

Expand All @@ -164,7 +164,7 @@
# as a global config setting for the whole program, or as a context manager
# (read its docs to learn more about caveats):

with tv_tensors.set_return_type("tv_tensor"):
with tv_tensors.set_return_type("TVTensor"):
new_bboxes = bboxes + 3
assert isinstance(new_bboxes, tv_tensors.BoundingBoxes)

Expand Down Expand Up @@ -203,17 +203,17 @@
# There are a few exceptions to this "unwrapping" rule:
# :meth:`~torch.Tensor.clone`, :meth:`~torch.Tensor.to`,
# :meth:`torch.Tensor.detach`, and :meth:`~torch.Tensor.requires_grad_` retain
# the tv_tensor type.
# the TVTensor type.
#
# Inplace operations on tv_tensors like ``obj.add_()`` will preserve the type of
# Inplace operations on TVTensors like ``obj.add_()`` will preserve the type of
# ``obj``. However, the **returned** value of inplace operations will be a pure
# tensor:

image = tv_tensors.Image([[[0, 1], [1, 0]]])

new_image = image.add_(1).mul_(2)

# image got transformed in-place and is still an Image tv_tensor, but new_image
# image got transformed in-place and is still a TVTensor Image, but new_image
# is a Tensor. They share the same underlying data and they're equal, just
# different classes.
assert isinstance(image, tv_tensors.Image)
Expand Down
Loading

0 comments on commit 25ec3f2

Please sign in to comment.