如何编写您自己的 TVTensor 类

注意

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

本指南适用于高级用户和下游库维护人员。我们将解释如何编写您自己的 TVTensor 类，以及如何使其与内置的 Torchvision v2 转换兼容。在继续之前，请确保您已阅读TVTensors 常见问题解答。

import torch
from torchvision import tv_tensors
from torchvision.transforms import v2

我们将创建一个非常简单的类，它仅从基类 TVTensor 类继承。它足以涵盖您需要了解的关于实现更精细的用例的内容。如果您需要创建一个携带元数据的类，请查看 BoundingBoxes 类的实现方式。

class MyTVTensor(tv_tensors.TVTensor):
    pass


my_dp = MyTVTensor([1, 2, 3])
my_dp

MyTVTensor([1., 2., 3.])

现在我们已经定义了自定义的 TVTensor 类，我们希望它与内置的 torchvision 转换和函数式 API 兼容。为此，我们需要实现一个内核，该内核执行转换的核心，然后通过 register_kernel() 将其“挂钩”到我们想要支持的函数。

我们在下面说明了这个过程：我们为 MyTVTensor 类的“水平翻转”操作创建了一个内核，并将其注册到函数式 API。

from torchvision.transforms.v2 import functional as F


@F.register_kernel(functional="hflip", tv_tensor_cls=MyTVTensor)
def hflip_my_tv_tensor(my_dp, *args, **kwargs):
    print("Flipping!")
    out = my_dp.flip(-1)
    return tv_tensors.wrap(out, like=my_dp)

要了解为什么使用 wrap()，请参阅我有一个 TVTensor，但现在我有一个 Tensor。怎么办！。现在忽略 *args, **kwargs，我们将在下面的参数转发以及确保内核的未来兼容性中进行解释。

注意

在我们上面对 register_kernel 的调用中，我们使用字符串 functional="hflip" 来引用我们要挂钩的函数。我们也可以使用函数本身，即 @register_kernel(functional=F.hflip, ...)。

现在我们已经注册了内核，我们可以在 MyTVTensor 实例上调用函数式 API

my_dp = MyTVTensor(torch.rand(3, 256, 256))
_ = F.hflip(my_dp)

Flipping!

我们也可以使用 RandomHorizontalFlip 转换，因为它在内部依赖于 hflip()

t = v2.RandomHorizontalFlip(p=1)
_ = t(my_dp)

Flipping!

注意

我们不能为转换类注册内核，我们只能为函数注册内核。我们不能注册转换类的原因是，一个转换可能在内部依赖于多个函数，因此通常我们不能为一个给定的类注册一个内核。

参数转发以及确保内核的未来兼容性

您要挂钩的函数式 API 是公开的，因此向后兼容：我们保证这些函数的参数不会在没有适当弃用周期的情况下被删除或重命名。但是，我们不保证向前兼容性，并且将来可能会添加新参数。

想象一下，在未来的版本中，Torchvision 向其 hflip() 函数添加了一个新的 inplace 参数。如果您已经定义并注册了自己的内核，如下所示

def hflip_my_tv_tensor(my_dp):  # noqa
    print("Flipping!")
    out = my_dp.flip(-1)
    return tv_tensors.wrap(out, like=my_dp)

那么调用 F.hflip(my_dp) 将会失败，因为 hflip 将尝试将新的 inplace 参数传递给您的内核，但您的内核不接受它。

因此，我们建议始终在内核的签名中使用 *args, **kwargs 来定义您的内核，如上所述。这样，您的内核将能够接受我们将来可能添加的任何新参数。（从技术上讲，仅添加 **kwargs 应该就足够了）。