TVTensors FAQ

注意

尝试在 Colab 上运行，或 滚动到末尾 下载完整的示例代码。

TVTensors 是与 `torchvision.transforms.v2` 一起引入的 Tensor 子类。本示例展示了这些 TVTensors 是什么以及它们的行为方式。

警告

目标读者 除非您正在编写自己的 transforms 或自己的 TVTensors，否则您可能不需要阅读本指南。这是一个相当底层的主题，大多数用户不需要担心：您不需要理解 TVTensors 的内部机制就可以有效地依赖 `torchvision.transforms.v2`。但它可能对尝试实现自己的数据集、transforms 或直接处理 TVTensors 的高级用户有用。

import PIL.Image

import torch
from torchvision import tv_tensors

什么是 TVTensors？

TVTensors 是零拷贝 Tensor 子类。

tensor = torch.rand(3, 256, 256)
image = tv_tensors.Image(tensor)

assert isinstance(image, torch.Tensor)
assert image.data_ptr() == tensor.data_ptr()

在底层，`torchvision.transforms.v2` 需要它们来正确地将输入数据分派到相应的函数。

`torchvision.tv_tensors` 支持五种 TVTensors 类型。

• 图像

• 视频

• 边界框

• 关键点

• 掩码

我能用 TVTensor 做什么？

TVTensors 看起来和感觉就像普通的 tensors — 它们 **就是** tensors。所有在普通 `torch.Tensor` 上支持的操作，如 `.sum()` 或任何 `torch.*` 运算符，同样适用于 TVTensors。有关一些注意事项，请参阅 我有一个 TVTensor，但现在我得到了一个 Tensor。怎么办？。

如何构建 TVTensor？

使用构造函数

每个 TVTensor 类都接受任何可以转换为 `Tensor` 的 tensor 类数据。

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

Image([[[[0, 1],
         [1, 0]]]], )

与其他的 PyTorch 创建操作类似，构造函数也接受 `dtype`、`device` 和 `requires_grad` 参数。

float_image = tv_tensors.Image([[[0, 1], [1, 0]]], dtype=torch.float32, requires_grad=True)
print(float_image)

Image([[[0., 1.],
        [1., 0.]]], grad_fn=<AliasBackward0>, )

此外，`Image` 和 `Mask` 也可以直接接受 `PIL.Image.Image`。

image = tv_tensors.Image(PIL.Image.open("../assets/astronaut.jpg"))
print(image.shape, image.dtype)

torch.Size([3, 512, 512]) torch.uint8

某些 TVTensors 需要传入额外的元数据才能构建。例如，`BoundingBoxes` 需要坐标格式以及对应图像的大小（`canvas_size`）以及实际值。这些元数据对于正确变换边界框是必需的。类似地，`KeyPoints` 也需要添加 `canvas_size` 元数据。

bboxes = tv_tensors.BoundingBoxes(
    [[17, 16, 344, 495], [0, 10, 0, 10]],
    format=tv_tensors.BoundingBoxFormat.XYXY,
    canvas_size=image.shape[-2:]
)
print(bboxes)


keypoints = tv_tensors.KeyPoints(
    [[17, 16], [344, 495], [0, 10], [0, 10]],
    canvas_size=image.shape[-2:]
)
print(keypoints)

BoundingBoxes([[ 17,  16, 344, 495],
               [  0,  10,   0,  10]], format=BoundingBoxFormat.XYXY, canvas_size=torch.Size([512, 512]), clamping_mode=soft)
KeyPoints([[ 17,  16],
           [344, 495],
           [  0,  10],
           [  0,  10]], canvas_size=torch.Size([512, 512]))

使用 `tv_tensors.wrap()`

您还可以使用 `wrap()` 函数将 tensor 对象包装成 TVTensor。当您已经拥有所需类型的对象时，这很有用，这通常发生在编写 transforms 时：您只想像输入一样包装输出。

new_bboxes = torch.tensor([0, 20, 30, 40])
new_bboxes = tv_tensors.wrap(new_bboxes, like=bboxes)
assert isinstance(new_bboxes, tv_tensors.BoundingBoxes)
assert new_bboxes.canvas_size == bboxes.canvas_size

`new_bboxes` 的元数据与 `bboxes` 相同，但您可以将其作为参数传入以覆盖它。

我有一个 TVTensor，但现在我得到了一个 Tensor。怎么办！

默认情况下，对 `TVTensor` 对象的操作将返回一个纯 Tensor。

assert isinstance(bboxes, tv_tensors.BoundingBoxes)

# Shift bboxes by 3 pixels in both H and W
new_bboxes = bboxes + 3

assert isinstance(new_bboxes, torch.Tensor)
assert not isinstance(new_bboxes, tv_tensors.BoundingBoxes)

注意

此行为仅影响原生的 `torch` 操作。如果您使用的是内置的 `torchvision` transforms 或 functionals，您将始终获得与输入相同类型的输出（纯 `Tensor` 或 `TVTensor`）。

但我想要一个 TVTensor 回来！

您可以将纯 tensor 重新包装成 TVTensor，只需调用 TVTensor 构造函数，或者使用 `wrap()` 函数（请参阅上方 如何构建 TVTensor？ 中的更多详细信息）。

new_bboxes = bboxes + 3
new_bboxes = tv_tensors.wrap(new_bboxes, like=bboxes)
assert isinstance(new_bboxes, tv_tensors.BoundingBoxes)

或者，您可以使用 `set_return_type()` 作为整个程序的全局配置设置，或者作为上下文管理器（请参阅其文档以了解更多关于注意事项的信息）。

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

为什么会发生这种情况？

出于性能原因。 `TVTensor` 类是 Tensor 子类，因此任何涉及 `TVTensor` 对象的运算都会通过 `__torch_function__` 协议。这会带来少量的开销，我们希望在可能的情况下避免这种开销。这对于内置的 `torchvision` transforms 来说无关紧要，因为我们可以在那里避免开销，但对于您模型中的 `forward` 来说，这可能是一个问题。

另一方面，也不是好太多。 对于保留 `TVTensor` 类型有意义的每种运算，都有同样多的运算返回纯 Tensor 是更可取的：例如，`img.sum()` 还是 `Image` 吗？如果我们一直保留 `TVTensor` 类型，那么即使是模型的 logits 或损失函数的输出也将是 `Image` 类型，这肯定是不理想的。

注意

这种行为是我们正在积极寻求反馈的。如果您觉得这令人惊讶，或者您对如何更好地支持您的用例有任何建议，请通过此 issue 与我们联系：https://github.com/pytorch/vision/issues/7319

例外

这个“解包”规则有几个例外：`clone()`、`to()`、`torch.Tensor.detach()` 和 `requires_grad_()` 保留 TVTensor 类型。

TVTensors 上的原地操作（如 `obj.add_()`）将保留 `obj` 的类型。但是，原地操作的 **返回** 值将是一个纯 tensor。

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

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

# 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)
print(image)

assert isinstance(new_image, torch.Tensor) and not isinstance(new_image, tv_tensors.Image)
assert (new_image == image).all()
assert new_image.data_ptr() == image.data_ptr()

Image([[[2, 4],
        [4, 2]]], )