如何编写自己的 v2 变换¶
本指南介绍如何编写与 torchvision 变换 V2 API 兼容的变换。
import torch
from torchvision import tv_tensors
from torchvision.transforms import v2
只需创建一个 nn.Module
并覆盖 forward
方法¶
在大多数情况下,只要您已经知道变换将期望的输入结构,这将是您需要做的所有事情。例如,如果您只是进行图像分类,您的变换通常会接受单个图像作为输入,或 (img, label)
输入。因此,您可以简单地将您的 forward
方法硬编码为仅接受该输入,例如
class MyCustomTransform(torch.nn.Module):
def forward(self, img, label):
# Do some transformations
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],
[136, 0, 173, 0]], format=BoundingBoxFormat.XYXY, canvas_size=(224, 224))
out_label = 3
注意
在代码中使用 TVTensor 类时,请务必熟悉以下部分:我有一个 TVTensor,但现在我有一个 Tensor。求助!
支持任意输入结构¶
在上一节中,我们假设您已经知道输入的结构,并且您对在代码中硬编码此预期结构感到满意。如果您希望您的自定义变换尽可能灵活,这可能会有些限制。
内置 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 都是张量子类)加上此处未计分的某些自定义逻辑 - 检查代码以获取详细信息。然后重新打包并返回(可能已转换的)条目,其结构与输入相同。
我们目前不提供面向开发人员的公共工具来实现这一点,但如果您认为这很有价值,请通过在我们的 GitHub 存储库 上创建一个问题来告知我们。
脚本的总运行时间:(0 分钟 0.006 秒)