merge_tensordicts¶
- class tensordict.merge_tensordicts(*tensordicts: T, callback_exist: Optional[Union[Callable[[Any], Any], Dict[NestedKey, Callable[[Any], Any]]]] = None)¶
合并 tensordict。
- 参数:
*tensordicts (TensorDict 或等效对象的序列) – 要合并的 tensordict 列表。
- 关键字参数:
callback_exist (可调用对象 或 Dict[str, 可调用对象], 可选) – 当每个 tensordict 中都存在某个条目时使用的可调用对象。如果条目存在于部分 tensordict 中但并非全部,或者如果
callback_exist
未传递,则使用 update 方法,并使用 tensordict 序列中的第一个非None
值。如果传递了一个可调用对象字典,它将包含传递给函数的一些嵌套键的关联回调函数。
示例
>>> from tensordict import merge_tensordicts, TensorDict >>> td0 = TensorDict({"a": {"b0": 0}, "c": {"d": {"e": 0}}, "common": 0}) >>> td1 = TensorDict({"a": {"b1": 1}, "f": {"g": {"h": 1}}, "common": 1}) >>> td2 = TensorDict({"a": {"b2": 2}, "f": {"g": {"h": 2}}, "common": 2}) >>> td = merge_tensordicts(td0, td1, td2, callback_exist=lambda *v: torch.stack(list(v))) >>> print(td) TensorDict( fields={ a: TensorDict( fields={ b0: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False), b1: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False), b2: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False), c: TensorDict( fields={ d: TensorDict( fields={ e: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False), common: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.int64, is_shared=False), f: TensorDict( fields={ g: TensorDict( fields={ h: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False) >>> print(td["common"]) tensor([0, 1, 2])