如何定制 v2 转换器

如何定制 v2 转换器#

本指南解释了如何编写与torchvision转换V2 API兼容的转换器。

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

只需创建 torch.nn.Module 并重写 forward 方法#

在大多数情况下，只要你已经知道你的转换将接受的输入结构，这就是你所需要的全部。例如，如果你只是进行图像分类，你的转换通常会接受单个图像作为输入，或者 (img, label) 输入。因此，你可以直接在你的 forward 方法中硬编码这些输入，例如：

class MyCustomTransform(torch.nn.Module):
    def forward(self, img, label):
        # 做一些变换
        return new_img, new_label

备注

这意味着，如果你有一个已经与V1转换（即 torchvision.transforms 中的那些）兼容的自定义转换，那么它在使用V2转换时仍然可以正常工作，无需任何更改！

我们将在下面更详细地说明这一点，以典型的检测案例为例，其中我们的样本只是图像、边界框和标签：

class MyCustomTransform(torch.nn.Module):
    def forward(self, img, bboxes, label):  # we assume inputs are always structured like this
        print(
            f"I'm transforming an image of shape {img.shape} "
            f"with bboxes = {bboxes}\n{label = }"
        )
        # Do some transformations. Here, we're just passing though the input
        return img, bboxes, label


transforms = v2.Compose([
    MyCustomTransform(),
    v2.RandomResizedCrop((224, 224), antialias=True),
    v2.RandomHorizontalFlip(p=1),
    v2.Normalize(mean=[0, 0, 0], std=[1, 1, 1])
])

H, W = 256, 256
img = torch.rand(3, H, W)
bboxes = tv_tensors.BoundingBoxes(
    torch.tensor([[0, 10, 10, 20], [50, 50, 70, 70]]),
    format="XYXY",
    canvas_size=(H, W)
)
label = 3

out_img, out_bboxes, out_label = transforms(img, bboxes, label)

I'm transforming an image of shape torch.Size([3, 256, 256]) with bboxes = BoundingBoxes([[ 0, 10, 10, 20],
               [50, 50, 70, 70]], format=BoundingBoxFormat.XYXY, canvas_size=(256, 256))
label = 3

print(f"Output image shape: {out_img.shape}\nout_bboxes = {out_bboxes}\n{out_label = }")

Output image shape: torch.Size([3, 224, 224])
out_bboxes = BoundingBoxes([[224,   0, 224,   0],
               [162,  23, 187,  44]], format=BoundingBoxFormat.XYXY, canvas_size=(224, 224))
out_label = 3

支持任意输入结构#

在上面的部分中，我们假设您已经了解输入的结构，并且可以在代码中硬编码这种预期结构。如果您希望自定义转换尽可能灵活，这可能会受到限制。

Torchvision V2内置转换的关键特性是它们可以接受任意输入结构，并返回相同结构的输出（带有转换后的元素）。例如，转换可以接收单个图像，或包含 (img, label) 的元组，或者任意嵌套字典作为输入：

structured_input = {
    "img": img,
    "annotations": (bboxes, label),
    "something_that_will_be_ignored": (1, "hello")
}
structured_output = v2.RandomHorizontalFlip(p=1)(structured_input)

assert isinstance(structured_output, dict)
assert structured_output["something_that_will_be_ignored"] == (1, "hello")
print(f"The transformed bboxes are:\n{structured_output['annotations'][0]}")

The transformed bboxes are:
BoundingBoxes([[246,  10, 256,  20],
               [186,  50, 206,  70]], format=BoundingBoxFormat.XYXY, canvas_size=(256, 256))

如果您想在您自己的转换中重现这种行为，我们邀请您查看我们的代码并根据需要进行调整。

简而言之，核心逻辑是将输入解包成一个平面列表，使用pytree，然后只转换那些可以转换的条目（决定是基于条目的类，因为所有TVTensor都是tensor子类），加上一些此处未详细说明的自定义逻辑 - 请查看代码以获取详细信息。然后，将（可能已转换的）条目重新打包并以与输入相同的结构返回。

目前，我们没有提供公开的开发工具来实现这一点，但如果这对您有价值，请通过在我们的GitHub repo上打开一个问题来告知我们。