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