快捷方式

ToTensorImage

class torchrl.envs.transforms.ToTensorImage(from_int: Optional[bool] = None, unsqueeze: bool = False, dtype: Optional[torch.device] = None, *, in_keys: Sequence[NestedKey] | None = None, out_keys: Sequence[NestedKey] | None = None, shape_tolerant: bool = False)[source]

将类似 numpy 的图像 (W x H x C) 转换为 pytorch 图像 (C x W x H)。

将观察图像从 (... x W x H x C) 张量转换为 (... x C x W x H) 张量。可选地,将输入张量从 [0, 255] 范围缩放到 [0.0, 1.0] 范围(详见 from_int)。

在其他情况下,张量不进行缩放直接返回。

参数:
  • from_int (bool, optional) – 如果为 True,则将张量从 [0, 255] 范围缩放到 [0.0, 1.0] 范围。如果为 False,则不缩放张量。如果为 None,则在张量不是浮点型张量时进行缩放。默认值为 None。

  • unsqueeze (bool) – 如果为 True,则沿第一维对观察张量进行 unsqueeze 操作。默认值为 False。

  • dtype (torch.dtype, optional) – 用于结果观察值的 dtype。

关键字参数:
  • in_keys (NestedKeys 列表) – 要处理的键。

  • out_keys (NestedKeys 列表) – 要写入的键。

  • shape_tolerant (bool, optional) – 如果为 True,将检查输入图像的形状。如果最后一个通道不是 3,则忽略置换。默认为 False

示例

>>> transform = ToTensorImage(in_keys=["pixels"])
>>> ri = torch.randint(0, 255, (1 , 1, 10, 11, 3), dtype=torch.uint8)
>>> td = TensorDict(
...     {"pixels": ri},
...     [1, 1])
>>> _ = transform(td)
>>> obs = td.get("pixels")
>>> print(obs.shape, obs.dtype)
torch.Size([1, 1, 3, 10, 11]) torch.float32
transform_observation_spec(observation_spec: TensorSpec) TensorSpec[source]

转换观察 spec,使结果 spec 与转换映射匹配。

参数:

observation_spec (TensorSpec) – 转换前的 spec

返回值:

转换后预期的 spec


© 版权所有 2022, Meta。

使用 Sphinx 构建,主题由 Read the Docs 提供。

文档

访问 PyTorch 的综合开发者文档

查看文档

教程

获取面向初学者和高级开发者的深度教程

查看教程

资源

查找开发资源并获得问题解答

查看资源