TVTensors 常见问题解答

注意

在 Colab 上尝试或 跳转至末尾 以下载完整的示例代码。

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

警告

目标读者 除非您正在编写自己的变换或 TVTensors，否则您可能不需要阅读本指南。这是一个相当底层的议题，大多数用户无需担心：您无需理解 TVTensors 的内部机制即可有效地使用 torchvision.transforms.v2。然而，对于尝试实现自己的数据集、变换或直接使用 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：

• Image

• Video

• BoundingBoxes

• Mask

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）以及实际值。这些元数据是正确变换边界框所必需的。

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

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

使用 tv_tensors.wrap()

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

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 变换或函数，您将始终获得与输入（纯 Tensor 或 TVTensor）相同的输出类型。

但我想要回 TVTensor！

您可以通过调用 TVTensor 构造函数，或者使用 wrap() 函数将纯 tensor 重新包装成 TVTensor（详见上文的 如何构造 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 变换而言，这并不重要，因为我们可以在那里避免开销，但在模型的 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]]], )