Shortcuts

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 = 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 的外观和感觉就像普通的张量 - 它们就是张量。 所有在普通 torch.Tensor 上支持的操作,比如 .sum() 或 任何 torch.* 操作符,也可以在 TVTensors 上使用。请参阅 我有一个 TVTensor 但现在我有一个 Tensor。救命! 了解一些注意事项。

如何构建一个TVTensor?

使用构造函数

每个TVTensor类都可以接受任何可以转换为Tensor的张量类数据

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

与其他PyTorch创建操作类似,构造函数也接受dtypedevicerequires_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>, )

此外,ImageMask 也可以直接接受一个 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()函数将一个张量对象包装成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转换或功能,您将始终获得与输入类型相同的输出(纯TensorTVTensor)。

但我想要一个TVTensor返回!

你可以通过调用TVTensor构造函数,或者使用wrap()函数(更多详情请参见如何构造一个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类型,这显然是不可取的。

注意

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

异常

这个“解包”规则有一些例外: clone(), to(), torch.Tensor.detach(), 和 requires_grad_() 保留 TVTensor 类型。

在TVTensors上的就地操作,如obj.add_(),将保留obj的类型。然而,就地操作的返回值将是一个纯张量:

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]]], )

脚本总运行时间: (0 分钟 0.009 秒)

Gallery generated by Sphinx-Gallery