入门 transforms v2¶

注意

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

此示例说明了您需要了解的关于新 torchvision.transforms.v2 API 的所有内容。我们将介绍图像分类等简单任务，以及对象检测/分割等更高级的任务。

首先，做一些准备工作

from pathlib import Path
import torch
import matplotlib.pyplot as plt
plt.rcParams["savefig.bbox"] = 'tight'

from torchvision.transforms import v2
from torchvision.io import decode_image

torch.manual_seed(1)

# If you're trying to run that on Colab, you can download the assets and the
# helpers from https://github.com/pytorch/vision/tree/main/gallery/
from helpers import plot
img = decode_image(str(Path('../assets') / 'astronaut.jpg'))
print(f"{type(img) = }, {img.dtype = }, {img.shape = }")

type(img) = <class 'torch.Tensor'>, img.dtype = torch.uint8, img.shape = torch.Size([3, 512, 512])

基础知识¶

Torchvision transforms 的行为类似于常规的 torch.nn.Module（事实上，它们中的大多数都是如此）：实例化一个 transform，传入一个输入，得到一个转换后的输出。

transform = v2.RandomCrop(size=(224, 224))
out = transform(img)

plot([img, out])

我只想做图像分类¶

如果您只关心图像分类，那么事情会非常简单。一个基本的分类流程可能如下所示：

transforms = v2.Compose([
    v2.RandomResizedCrop(size=(224, 224), antialias=True),
    v2.RandomHorizontalFlip(p=0.5),
    v2.ToDtype(torch.float32, scale=True),
    v2.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
out = transforms(img)

plot([img, out])

这种转换流程通常作为 transform 参数传递给 Datasets，例如 ImageNet(..., transform=transforms)。

差不多就是这样了。接下来，请阅读我们的主文档，了解更多关于推荐实践和约定，或者探索更多示例，例如如何使用增强转换，如 CutMix 和 MixUp。

注意

如果您已经依赖 torchvision.transforms v1 API，我们建议您切换到新的 v2 transforms。这非常简单：v2 transforms 完全兼容 v1 API，所以您只需要更改导入即可！

视频、边界框、掩码、关键点¶

来自 torchvision.transforms.v2 命名空间中的 Torchvision transforms 支持图像分类以外的任务：它们还可以转换旋转或轴对齐的边界框、分割/检测掩码、视频和关键点。

让我们简要看一个带有边界框的检测示例。

from torchvision import tv_tensors  # we'll describe this a bit later, bare with us

boxes = tv_tensors.BoundingBoxes(
    [
        [15, 10, 370, 510],
        [275, 340, 510, 510],
        [130, 345, 210, 425]
    ],
    format="XYXY", canvas_size=img.shape[-2:])

transforms = v2.Compose([
    v2.RandomResizedCrop(size=(224, 224), antialias=True),
    v2.RandomPhotometricDistort(p=1),
    v2.RandomHorizontalFlip(p=1),
])
out_img, out_boxes = transforms(img, boxes)
print(type(boxes), type(out_boxes))

plot([(img, boxes), (out_img, out_boxes)])

<class 'torchvision.tv_tensors._bounding_boxes.BoundingBoxes'> <class 'torchvision.tv_tensors._bounding_boxes.BoundingBoxes'>

上面的示例侧重于对象检测。但如果我们有用于对象分割或语义分割的掩码（torchvision.tv_tensors.Mask）或视频（torchvision.tv_tensors.Video），我们可以以完全相同的方式将它们传递给 transforms。

现在您可能还有几个问题：什么是 TVTensors？如何使用它们？这些 transforms 的预期输入/输出是什么？我们将在接下来的部分回答这些问题。

什么是 TVTensors？¶

TVTensors 是 torch.Tensor 的子类。可用的 TVTensors 有 Image、BoundingBoxes、Mask、Video 和 KeyPoints。

TVTensors 的外观和感觉与普通张量一样——它们就是张量。在普通 torch.Tensor 上支持的所有操作，例如 .sum() 或任何 torch.* 运算符，同样适用于 TVTensors。

img_dp = tv_tensors.Image(torch.randint(0, 256, (3, 256, 256), dtype=torch.uint8))

print(f"{isinstance(img_dp, torch.Tensor) = }")
print(f"{img_dp.dtype = }, {img_dp.shape = }, {img_dp.sum() = }")

isinstance(img_dp, torch.Tensor) = True
img_dp.dtype = torch.uint8, img_dp.shape = torch.Size([3, 256, 256]), img_dp.sum() = tensor(25087958)

这些 TVTensor 类是 transforms 的核心：为了转换给定的输入，transforms 首先查看对象的类，然后相应地分派到适当的实现。

您目前不需要了解更多关于 TVTensors 的知识，但想要了解更多的高级用户可以参考 TVTensors FAQ。

我应该传入什么作为输入？¶

上面我们已经看到了两个示例：一个示例中我们传入了单个图像作为输入，即 out = transforms(img)；另一个示例中我们同时传入了图像和边界框，即 out_img, out_boxes = transforms(img, boxes)。

事实上，transforms 支持任意输入结构。输入可以是单个图像、元组、任意嵌套的字典……几乎任何东西。输出将返回相同的结构。下面，我们使用相同的检测 transforms，但将元组（图像，目标字典）作为输入，并获得了相同的输出结构。

target = {
    "boxes": boxes,
    "labels": torch.arange(boxes.shape[0]),
    "this_is_ignored": ("arbitrary", {"structure": "!"})
}

# Re-using the transforms and definitions from above.
out_img, out_target = transforms(img, target)

plot([(img, target["boxes"]), (out_img, out_target["boxes"])])
print(f"{out_target['this_is_ignored']}")

('arbitrary', {'structure': '!'})

我们传入了一个元组，所以我们得到一个元组返回，第二个元素是转换后的目标字典。Transforms 并不真正关心输入的结构；正如上面提到的，它们只关心对象的类型并相应地转换它们。

*外部*对象，如字符串或整数，将被简单地传递。这可能很有用，例如，如果您想在调试时将路径与每个样本关联起来！

注意

免责声明：此说明稍微高级，初次阅读时可以安全跳过。

纯粹的 torch.Tensor 对象通常被视为图像（或对于特定于视频的 transforms，被视为视频）。事实上，您可能已经注意到，在上面的代码中，我们根本没有使用 Image 类，但我们的图像仍然得到了正确转换。Transforms 遵循以下逻辑来确定一个纯张量是应被视为图像（或视频），还是仅被忽略：

如果输入中存在 Image、Video 或 PIL.Image.Image 实例，所有其他纯张量都将被传递。
如果不存在 Image 或 Video 实例，只有第一个纯 torch.Tensor 将被转换为图像或视频，而所有其他张量将被传递。这里的“第一个”意味着“在深度优先遍历中的第一个”。

这正是上面检测示例中所发生的情况：第一个纯张量是图像，因此它被正确转换，而所有其他纯张量实例，如 labels，都被传递（尽管标签仍然可以被某些 transforms 转换，如 SanitizeBoundingBoxes！）。

Transforms 和 Datasets 的互操作性¶

粗略地说，datasets 的输出必须对应于 transforms 的输入。如何实现这一点取决于您使用的是 torchvision 的内置数据集，还是您自己的自定义数据集。

使用内置数据集¶

如果您只是做图像分类，您不需要做任何事情。只需使用数据集的 transform 参数，例如 ImageNet(..., transform=transforms)，您就可以开始使用了。

Torchvision 还支持对象检测或分割的数据集，如 torchvision.datasets.CocoDetection。这些数据集在 torchvision.transforms.v2 模块和 TVTensors 出现之前就已存在，因此它们默认不会返回 TVTensors。

一个强制这些数据集返回 TVTensors 并使其与 v2 transforms 兼容的简单方法是使用 torchvision.datasets.wrap_dataset_for_transforms_v2() 函数。

from torchvision.datasets import CocoDetection, wrap_dataset_for_transforms_v2

dataset = CocoDetection(..., transforms=my_transforms)
dataset = wrap_dataset_for_transforms_v2(dataset)
# Now the dataset returns TVTensors!

使用您自己的数据集¶

如果您有自定义数据集，那么您需要将您的对象转换为相应的 TVTensor 类。创建 TVTensor 实例非常简单，有关更多详细信息，请参阅如何构造一个 TVTensor？。

您可以在以下两个主要地方实现该转换逻辑：

在数据集的 __getitem__ 方法的末尾，在返回样本之前（或通过子类化数据集）。
作为您的 transforms 管道的第一步。

无论哪种方式，逻辑都将取决于您的具体数据集。

脚本总运行时间：（0 分钟 0.837 秒）

由 Sphinx-Gallery 生成的画廊