QueryModule¶
- class torchrl.data.QueryModule(*args, **kwargs)[source]¶
用于生成兼容存储索引的模块。
该模块查询存储并返回该存储所需的索引。目前,它仅输出整数索引 (torch.int64)。
- 参数::
in_keys (NestedKeys 列表) – 输入 tensordict 的键,将用于生成哈希值。
index_key (NestedKey) – 将写入索引值的输出键。默认为
"_index"
。
- 关键字参数::
hash_key (NestedKey) – 将写入哈希值的输出键。默认为
"_hash"
。hash_module (可调用对象[[Any], int] 或此类对象的列表,可选) – 类似于
SipHash
的哈希模块(默认)。如果提供了可调用对象的列表,其长度必须与 in_keys 的数量相等。hash_to_int (可调用对象[[int], int],可选) – 一个有状态函数,将哈希值映射到存储中对应的非负整数索引。默认为
HashToInt
。aggregator (可调用对象[[int], int],可选) – 用于将多个哈希值组合在一起的哈希函数。仅当存在多个
in_keys
时才应传递此参数。如果提供了单个hash_module
但未传递聚合器,则其值将取 hash_module 的值。如果没有hash_module
或hash_module
的列表但未传递聚合器,则将默认为SipHash
。clone (布尔值,可选) – 如果为
True
,将返回输入 TensorDict 的浅层副本。这可用于检索存储中的整数索引,对应于给定的输入 tensordict。这可以在运行时通过提供clone
参数覆盖。默认为False
。
示例
>>> query_module = QueryModule( ... in_keys=["key1", "key2"], ... index_key="index", ... hash_module=SipHash(), ... ) >>> query = TensorDict( ... { ... "key1": torch.Tensor([[1], [1], [1], [2]]), ... "key2": torch.Tensor([[3], [3], [2], [3]]), ... "other": torch.randn(4), ... }, ... batch_size=(4,), ... ) >>> res = query_module(query) >>> # The first two pairs of key1 and key2 match >>> assert res["index"][0] == res["index"][1] >>> # The last three pairs of key1 and key2 have at least one mismatching value >>> assert res["index"][1] != res["index"][2] >>> assert res["index"][2] != res["index"][3]
- add_module(name: str, module: Optional[Module]) None ¶
向当前模块添加一个子模块。
可以使用给定的名称作为属性访问该模块。
- 参数::
name (字符串) – 子模块的名称。可以使用给定的名称从此模块访问子模块
module (模块) – 要添加到模块的子模块。
- apply(fn: Callable[[Module], None]) T ¶
将
fn
递归地应用于每个子模块 (由.children()
返回) 以及自身。典型用法包括初始化模型的参数(另请参见 torch.nn.init)。
- 参数::
fn (模块 -> None) – 要应用于每个子模块的函数
- 返回值::
自身
- 返回值类型::
模块
示例
>>> @torch.no_grad() >>> def init_weights(m): >>> print(m) >>> if type(m) == nn.Linear: >>> m.weight.fill_(1.0) >>> print(m.weight) >>> net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2)) >>> net.apply(init_weights) Linear(in_features=2, out_features=2, bias=True) Parameter containing: tensor([[1., 1.], [1., 1.]], requires_grad=True) Linear(in_features=2, out_features=2, bias=True) Parameter containing: tensor([[1., 1.], [1., 1.]], requires_grad=True) Sequential( (0): Linear(in_features=2, out_features=2, bias=True) (1): Linear(in_features=2, out_features=2, bias=True) )
- bfloat16() T ¶
将所有浮点参数和缓冲区转换为
bfloat16
数据类型。注意
此方法会就地修改模块。
- 返回值::
自身
- 返回值类型::
模块
- buffers(recurse: bool = True) Iterator[Tensor] ¶
返回模块缓冲区的迭代器。
- 参数::
recurse (布尔值) – 如果为 True,则生成此模块及所有子模块的缓冲区。否则,仅生成作为此模块直接成员的缓冲区。
- 生成值::
torch.Tensor – 模块缓冲区
示例
>>> # xdoctest: +SKIP("undefined vars") >>> for buf in model.buffers(): >>> print(type(buf), buf.size()) <class 'torch.Tensor'> (20L,) <class 'torch.Tensor'> (20L, 1L, 5L, 5L)
- compile(*args, **kwargs)¶
使用
torch.compile()
编译此模块的 forward 方法。此模块的 __call__ 方法被编译,所有参数都按原样传递给
torch.compile()
。请参阅
torch.compile()
以了解有关此函数参数的详细信息。
- cpu() T ¶
将所有模型参数和缓冲区移动到 CPU。
注意
此方法会就地修改模块。
- 返回值::
自身
- 返回值类型::
模块
- cuda(device: Optional[Union[int, device]] = None) T ¶
将所有模型参数和缓冲区移动到 GPU。
这也会使相关的参数和缓冲区成为不同的对象。因此,如果模块在优化期间将驻留在 GPU 上,则应在构建优化器之前调用此方法。
注意
此方法会就地修改模块。
- 参数::
device (整数,可选) – 如果指定,所有参数将复制到该设备
- 返回值::
自身
- 返回值类型::
模块
- double() T ¶
将所有浮点参数和缓冲区转换为
double
数据类型。注意
此方法会就地修改模块。
- 返回值::
自身
- 返回值类型::
模块
- eval() T ¶
将模块设置为评估模式。
这仅对某些模块有效。有关特定模块在训练/评估模式下的行为详情,即它们是否受影响(例如
Dropout
,BatchNorm
等),请参阅其文档。这等同于
self.train(False)
。请参阅 本地禁用梯度计算 以比较 .eval() 以及其他一些可能与其混淆的类似机制。
- 返回值::
自身
- 返回值类型::
模块
- extra_repr() str ¶
返回模块的额外表示。
要打印定制的额外信息,应在自己的模块中重新实现此方法。单行和多行字符串均可接受。
- float() T ¶
将所有浮点参数和缓冲区转换为
float
数据类型。注意
此方法会就地修改模块。
- 返回值::
自身
- 返回值类型::
模块
- forward(tensordict: TensorDictBase, *, extend: bool = True, write_hash: bool = True, clone: bool | None = None) TensorDictBase [source]¶
定义每次调用时执行的计算。
所有子类都应覆盖此方法。
注意
虽然需要在函数内定义 forward pass 的实现,但应调用
Module
实例而不是此方法,因为前者会负责运行已注册的 hook,而后者会静默忽略它们。
- get_buffer(target: str) Tensor ¶
返回由
target
给定的缓冲区(如果存在),否则抛出错误。请参阅
get_submodule
的文档字符串,以了解此方法的更详细功能说明以及如何正确指定target
。- 参数::
target – 要查找的缓冲区的完全限定字符串名称。(请参阅
get_submodule
以了解如何指定完全限定字符串。)- 返回值::
由
target
引用的缓冲区- 返回值类型::
- 抛出::
AttributeError – 如果 target 字符串引用了无效路径或解析到的内容不是缓冲区
- get_extra_state() Any ¶
返回要包含在模块 state_dict 中的所有额外状态。
如果你的模块需要存储额外状态,则实现此方法以及对应的
set_extra_state()
。构建模块的 state_dict() 时会调用此函数。请注意,额外状态应该是可 pickle 化的,以确保 state_dict 的序列化正常工作。我们仅提供 Tensor 序列化的向后兼容性保证;其他对象如果其序列化后的 pickled 形式发生变化,可能会破坏向后兼容性。
- 返回值::
要存储在模块 state_dict 中的任何额外状态
- 返回值类型::
对象
- get_parameter(target: str) Parameter ¶
返回由
target
给定的参数(如果存在),否则抛出错误。请参阅
get_submodule
的文档字符串,以了解此方法的更详细功能说明以及如何正确指定target
。- 参数::
target – 要查找的 Parameter 的完全限定字符串名称。(请参阅
get_submodule
以了解如何指定完全限定字符串。)- 返回值::
由
target
引用的 Parameter- 返回值类型::
torch.nn.Parameter
- 抛出::
AttributeError – 如果 target 字符串引用了无效路径或解析到的内容不是
nn.Parameter
- get_submodule(target: str) Module ¶
返回由
target
给定的子模块(如果存在),否则抛出错误。例如,假设你有一个
nn.Module
A
,结构如下A( (net_b): Module( (net_c): Module( (conv): Conv2d(16, 33, kernel_size=(3, 3), stride=(2, 2)) ) (linear): Linear(in_features=100, out_features=200, bias=True) ) )
(图示为一个
nn.Module
A
。A
有一个嵌套子模块net_b
,而它自身有两个子模块net_c
和linear
。net_c
然后有一个子模块conv
。)要检查我们是否拥有
linear
子模块,我们会调用get_submodule("net_b.linear")
。要检查我们是否拥有conv
子模块,我们会调用get_submodule("net_b.net_c.conv")
。get_submodule
的运行时由target
中模块嵌套的层级决定。对named_modules
的查询可以达到相同的结果,但在传递性模块数量上是 O(N)。因此,对于检查某个子模块是否存在这样的简单任务,应始终使用get_submodule
。- 参数::
target – 要查找的子模块的完全限定字符串名称。(有关如何指定完全限定字符串,请参阅上述示例。)
- 返回值::
由
target
引用的子模块- 返回值类型::
- 抛出::
AttributeError – 如果 target 字符串引用了无效路径或解析到的内容不是
nn.Module
- half() T ¶
将所有浮点参数和缓冲区转换为
half
数据类型。注意
此方法会就地修改模块。
- 返回值::
自身
- 返回值类型::
模块
- ipu(device: Optional[Union[int, device]] = None) T ¶
将所有模型参数和缓冲区移动到 IPU。
这也会使相关的参数和缓冲区成为不同的对象。因此,如果模块在优化期间将驻留在 IPU 上,则应在构建优化器之前调用此方法。
注意
此方法会就地修改模块。
- 参数::
device (整数,可选) – 如果指定,所有参数将复制到该设备
- 返回值::
自身
- 返回值类型::
模块
- static is_tdmodule_compatible(module)¶
检查模块是否与 TensorDictModule API 兼容。
- load_state_dict(state_dict: Mapping[str, Any], strict: bool = True, assign: bool = False)¶
将参数和缓冲区从
state_dict
复制到此模块及其后代。如果
strict
为True
,则state_dict
的键必须与此模块的state_dict()
函数返回的键完全匹配。警告
如果
assign
为True
,则必须在调用load_state_dict
之后创建优化器,除非get_swap_module_params_on_conversion()
为True
。- 参数::
state_dict (字典) – 包含参数和持久缓冲区的字典。
strict (布尔值,可选) – 是否严格要求
state_dict
中的键与此模块的state_dict()
函数返回的键匹配。默认值:True
assign (bool, 可选) – 设置为
False
时,将保留当前模块中张量的属性,而设置为True
时,则保留状态字典 (state dict) 中张量的属性。唯一的例外是requires_grad
字段。默认值: ``False`
- 返回值::
- missing_keys 是一个字符串列表,包含此模块预期存在但在提供的
state_dict
中缺失的键。
- missing_keys 是一个字符串列表,包含此模块预期存在但在提供的
- unexpected_keys 是一个字符串列表,包含此模块不预期存在但在提供的
state_dict
中存在的键。
- unexpected_keys 是一个字符串列表,包含此模块不预期存在但在提供的
- 返回值类型::
NamedTuple
,包含missing_keys
和unexpected_keys
字段。
注意
如果一个参数或缓冲区被注册为
None
,并且其对应的键存在于state_dict
中,则load_state_dict()
将引发RuntimeError
。
- modules() Iterator[Module] ¶
返回一个迭代器,遍历网络中的所有模块。
- 生成值::
Module – 网络中的一个模块
注意
重复的模块只返回一次。在下面的例子中,
l
只会返回一次。示例
>>> l = nn.Linear(2, 2) >>> net = nn.Sequential(l, l) >>> for idx, m in enumerate(net.modules()): ... print(idx, '->', m) 0 -> Sequential( (0): Linear(in_features=2, out_features=2, bias=True) (1): Linear(in_features=2, out_features=2, bias=True) ) 1 -> Linear(in_features=2, out_features=2, bias=True)
- mtia(device: Optional[Union[int, device]] = None) T ¶
将所有模型参数和缓冲区移动到 MTIA。
这也会使相关的参数和缓冲区成为不同的对象。因此,如果在优化时模块将保留在 MTIA 上,则应在构建优化器之前调用此方法。
注意
此方法会就地修改模块。
- 参数::
device (整数,可选) – 如果指定,所有参数将复制到该设备
- 返回值::
自身
- 返回值类型::
模块
- named_buffers(prefix: str = '', recurse: bool = True, remove_duplicate: bool = True) Iterator[Tuple[str, Tensor]] ¶
返回一个迭代器,遍历模块缓冲区,同时产生缓冲区名称和缓冲区本身。
- 参数::
prefix (str) – 添加到所有缓冲区名称前的字符串前缀。
recurse (bool, 可选) – 如果为 True,则产生此模块及其所有子模块的缓冲区。否则,仅产生直接属于此模块的缓冲区。默认为 True。
remove_duplicate (bool, 可选) – 是否移除结果中重复的缓冲区。默认为 True。
- 生成值::
(str, torch.Tensor) – 包含名称和缓冲区的元组。
示例
>>> # xdoctest: +SKIP("undefined vars") >>> for name, buf in self.named_buffers(): >>> if name in ['running_var']: >>> print(buf.size())
- named_children() Iterator[Tuple[str, Module]] ¶
返回一个迭代器,遍历直接子模块,同时产生模块名称和模块本身。
- 生成值::
(str, Module) – 包含名称和子模块的元组。
示例
>>> # xdoctest: +SKIP("undefined vars") >>> for name, module in model.named_children(): >>> if name in ['conv4', 'conv5']: >>> print(module)
- named_modules(memo: Optional[Set[Module]] = None, prefix: str = '', remove_duplicate: bool = True)¶
返回一个迭代器,遍历网络中的所有模块,同时产生模块名称和模块本身。
- 参数::
memo – 用于存储已添加到结果中的模块集合的备忘录。
prefix – 将添加到模块名称中的前缀。
remove_duplicate – 是否移除结果中重复的模块实例。
- 生成值::
(str, Module) – 包含名称和模块的元组。
注意
重复的模块只返回一次。在下面的例子中,
l
只会返回一次。示例
>>> l = nn.Linear(2, 2) >>> net = nn.Sequential(l, l) >>> for idx, m in enumerate(net.named_modules()): ... print(idx, '->', m) 0 -> ('', Sequential( (0): Linear(in_features=2, out_features=2, bias=True) (1): Linear(in_features=2, out_features=2, bias=True) )) 1 -> ('0', Linear(in_features=2, out_features=2, bias=True))
- named_parameters(prefix: str = '', recurse: bool = True, remove_duplicate: bool = True) Iterator[Tuple[str, Parameter]] ¶
返回一个迭代器,遍历模块参数,同时产生参数名称和参数本身。
- 参数::
prefix (str) – 添加到所有参数名称前的字符串前缀。
recurse (bool) – 如果为 True,则产生此模块及其所有子模块的参数。否则,仅产生直接属于此模块的参数。
remove_duplicate (bool, 可选) – 是否移除结果中重复的参数。默认为 True。
- 生成值::
(str, Parameter) – 包含名称和参数的元组。
示例
>>> # xdoctest: +SKIP("undefined vars") >>> for name, param in self.named_parameters(): >>> if name in ['bias']: >>> print(param.size())
- parameters(recurse: bool = True) Iterator[Parameter] ¶
返回一个迭代器,遍历模块参数。
这通常传递给优化器。
- 参数::
recurse (bool) – 如果为 True,则产生此模块及其所有子模块的参数。否则,仅产生直接属于此模块的参数。
- 生成值::
Parameter – 模块参数
示例
>>> # xdoctest: +SKIP("undefined vars") >>> for param in model.parameters(): >>> print(type(param), param.size()) <class 'torch.Tensor'> (20L,) <class 'torch.Tensor'> (20L, 1L, 5L, 5L)
- register_backward_hook(hook: Callable[[Module, Union[Tuple[Tensor, ...], Tensor], Union[Tuple[Tensor, ...], Tensor], Union[None, Tuple[Tensor, ...], Tensor]]) RemovableHandle ¶
在模块上注册一个反向传播钩子。
此函数已弃用,请改用
register_full_backward_hook()
。此函数的行为在未来版本中将发生变化。- 返回值::
一个句柄,可以通过调用
handle.remove()
来移除添加的钩子。- 返回值类型::
torch.utils.hooks.RemovableHandle
- register_buffer(name: str, tensor: Optional[Tensor], persistent: bool = True) None ¶
向模块添加一个缓冲区。
这通常用于注册不应被视为模型参数的缓冲区。例如,BatchNorm 的
running_mean
不是参数,但它是模块状态的一部分。缓冲区默认是持久的,并将与参数一起保存。可以通过将persistent
设置为False
来改变此行为。持久缓冲区和非持久缓冲区的唯一区别在于后者不会成为此模块的state_dict
的一部分。可以使用给定的名称作为属性访问缓冲区。
- 参数::
name (str) – 缓冲区的名称。可以使用给定的名称从该模块访问缓冲区。
tensor (Tensor 或 None) – 要注册的缓冲区。如果为
None
,则对缓冲区执行的操作(如cuda
)将被忽略。如果为None
,则该缓冲区将不包含在模块的state_dict
中。persistent (bool) – 缓冲区是否为此模块
state_dict
的一部分。
示例
>>> # xdoctest: +SKIP("undefined vars") >>> self.register_buffer('running_mean', torch.zeros(num_features))
- register_forward_hook(hook: Union[Callable[[T, Tuple[Any, ...], Any], Optional[Any]], Callable[[T, Tuple[Any, ...], Dict[str, Any], Any], Optional[Any]]], *, prepend: bool = False, with_kwargs: bool = False, always_call: bool = False) RemovableHandle ¶
在模块上注册一个前向传播钩子。
每次
forward()
计算出输出后,都会调用此钩子。如果
with_kwargs
为False
或未指定,则输入仅包含传递给模块的位置参数。关键字参数不会传递给钩子,只会传递给forward
。钩子可以修改输出。它可以就地修改输入,但由于此钩子在forward()
调用后才被调用,因此修改输入不会影响前向传播。钩子应具有以下签名:hook(module, args, output) -> None or modified output
如果
with_kwargs
为True
,前向钩子将接收传递给 forward 函数的kwargs
,并期望返回可能修改过的输出。钩子应具有以下签名:hook(module, args, kwargs, output) -> None or modified output
- 参数::
hook (Callable) – 要注册的用户定义的钩子。
prepend (bool) – 如果为
True
,则提供的hook
将在此torch.nn.modules.Module
上的所有现有forward
钩子之前触发。否则,提供的hook
将在此torch.nn.modules.Module
上的所有现有forward
钩子之后触发。请注意,使用register_module_forward_hook()
注册的全局forward
钩子将在此方法注册的所有钩子之前触发。默认值:False
。with_kwargs (bool) – 如果为
True
,则hook
将接收传递给 forward 函数的 kwargs。默认值:False
。always_call (bool) – 如果为
True
,则无论调用模块时是否引发异常,都会运行hook
。默认值:False
。
- 返回值::
一个句柄,可以通过调用
handle.remove()
来移除添加的钩子。- 返回值类型::
torch.utils.hooks.RemovableHandle
- register_forward_pre_hook(hook: Union[Callable[[T, Tuple[Any, ...]], Optional[Any]], Callable[[T, Tuple[Any, ...], Dict[str, Any]], Optional[Tuple[Any, Dict[str, Any]]]]], *, prepend: bool = False, with_kwargs: bool = False) RemovableHandle ¶
在模块上注册一个前向传播预钩子。
每次调用
forward()
之前,都会调用此钩子。如果
with_kwargs
为 false 或未指定,则输入仅包含传递给模块的位置参数。关键字参数不会传递给钩子,只会传递给forward
。钩子可以修改输入。用户可以在钩子中返回一个元组或一个修改后的单个值。如果返回单个值(除非该值本身就是一个元组),我们将将其包装成一个元组。钩子应具有以下签名:hook(module, args) -> None or modified input
如果
with_kwargs
为 true,前向传播预钩子将接收传递给 forward 函数的 kwargs。如果钩子修改了输入,则应同时返回 args 和 kwargs。钩子应具有以下签名:hook(module, args, kwargs) -> None or a tuple of modified input and kwargs
- 参数::
hook (Callable) – 要注册的用户定义的钩子。
prepend (bool) – 如果为 true,则提供的
hook
将在此torch.nn.modules.Module
上的所有现有forward_pre
钩子之前触发。否则,提供的hook
将在此torch.nn.modules.Module
上的所有现有forward_pre
钩子之后触发。请注意,使用register_module_forward_pre_hook()
注册的全局forward_pre
钩子将在此方法注册的所有钩子之前触发。默认值:False
。with_kwargs (bool) – 如果为 true,则
hook
将接收传递给 forward 函数的 kwargs。默认值:False
。
- 返回值::
一个句柄,可以通过调用
handle.remove()
来移除添加的钩子。- 返回值类型::
torch.utils.hooks.RemovableHandle
- register_full_backward_hook(hook: Callable[[Module, Union[Tuple[Tensor, ...], Tensor], Union[Tuple[Tensor, ...], Tensor], Union[None, Tuple[Tensor, ...], Tensor]], prepend: bool = False) RemovableHandle ¶
在模块上注册一个反向传播钩子。
每次计算模块的梯度时,都会调用此钩子,即只有在计算模块输出的梯度时,此钩子才会执行。钩子应具有以下签名:
hook(module, grad_input, grad_output) -> tuple(Tensor) or None
grad_input
和grad_output
是分别包含输入和输出梯度的元组。钩子不应修改其参数,但可以选择返回一个新的输入梯度,该梯度将用于后续计算中替代grad_input
。grad_input
仅对应于作为位置参数给出的输入,所有 kwarg 参数都被忽略。grad_input
和grad_output
中的所有非 Tensor 参数的对应条目将为None
。出于技术原因,当此钩子应用于模块时,其 forward 函数将接收传递给模块的每个张量的视图。同样,调用者将接收模块 forward 函数返回的每个张量的视图。
警告
使用反向传播钩子时,不允许就地修改输入或输出,否则将引发错误。
- 参数::
hook (Callable) – 要注册的用户定义的钩子。
prepend (bool) – 如果为 true,则提供的
hook
将在此torch.nn.modules.Module
上的所有现有backward
钩子之前触发。否则,提供的hook
将在此torch.nn.modules.Module
上的所有现有backward
钩子之后触发。请注意,使用register_module_full_backward_hook()
注册的全局backward
钩子将在此方法注册的所有钩子之前触发。
- 返回值::
一个句柄,可以通过调用
handle.remove()
来移除添加的钩子。- 返回值类型::
torch.utils.hooks.RemovableHandle
- register_full_backward_pre_hook(hook: Callable[[Module, Union[Tuple[Tensor, ...], Tensor]], Union[None, Tuple[Tensor, ...], Tensor]], prepend: bool = False) RemovableHandle ¶
在模块上注册一个反向传播预钩子。
每次计算模块的梯度时,都会调用此钩子。钩子应具有以下签名:
hook(module, grad_output) -> tuple[Tensor] or None
grad_output
是一个元组。该 hook 不应修改其参数,但可以选择返回一个相对于输出的新梯度,此新梯度将替代grad_output
用于后续计算。grad_output
中所有非 Tensor 参数的条目都将为None
。出于技术原因,当此钩子应用于模块时,其 forward 函数将接收传递给模块的每个张量的视图。同样,调用者将接收模块 forward 函数返回的每个张量的视图。
警告
使用反向 hook 时不允许就地修改输入,否则将引发错误。
- 参数::
hook (Callable) – 要注册的用户定义的钩子。
prepend (bool) – 如果为 True,则提供的
hook
将在此torch.nn.modules.Module
上所有现有backward_pre
hook 之前触发。否则,提供的hook
将在此torch.nn.modules.Module
上所有现有backward_pre
hook 之后触发。请注意,使用register_module_full_backward_pre_hook()
注册的全局backward_pre
hook 将在此方法注册的所有 hook 之前触发。
- 返回值::
一个句柄,可以通过调用
handle.remove()
来移除添加的钩子。- 返回值类型::
torch.utils.hooks.RemovableHandle
- register_load_state_dict_post_hook(hook)¶
注册一个后置 hook,在模块的
load_state_dict()
被调用后运行。- 它应具有以下签名:
hook(module, incompatible_keys) -> None
参数
module
是注册此 hook 的当前模块,参数incompatible_keys
是一个NamedTuple
,包含属性missing_keys
和unexpected_keys
。missing_keys
是一个包含缺失键的str
list
,unexpected_keys
是一个包含意外键的str
list
。如果需要,可以就地修改给定的 incompatible_keys。
请注意,调用
load_state_dict()
时使用strict=True
进行的检查会受到 hook 对missing_keys
或unexpected_keys
所做修改的影响,这是符合预期的。向任何一组键中添加内容都会在strict=True
时导致抛出错误,而清空缺失和意外键则可以避免错误。- 返回值::
一个句柄,可以通过调用
handle.remove()
来移除添加的钩子。- 返回值类型::
torch.utils.hooks.RemovableHandle
- register_load_state_dict_pre_hook(hook)¶
注册一个前置 hook,在模块的
load_state_dict()
被调用前运行。- 它应具有以下签名:
hook(module, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs) -> None # noqa: B950
- 参数::
hook (Callable) – 可调用 hook,将在加载状态字典前被调用。
- register_module(name: str, module: Optional[Module]) None ¶
add_module()
的别名。
- register_parameter(name: str, param: Optional[Parameter]) None ¶
向模块添加一个参数。
该参数可以使用给定名称作为属性访问。
- 参数::
name (str) – 参数的名称。可以使用给定名称从此模块访问该参数
param (Parameter 或 None) – 要添加到模块的参数。如果为
None
,则 `cuda` 等在参数上运行的操作将被忽略。如果为None
,则该参数将不包含在模块的state_dict
中。
- register_state_dict_post_hook(hook)¶
为
state_dict()
方法注册一个后置 hook。- 它应具有以下签名:
hook(module, state_dict, prefix, local_metadata) -> None
注册的 hook 可以就地修改
state_dict
。
- register_state_dict_pre_hook(hook)¶
为
state_dict()
方法注册一个前置 hook。- 它应具有以下签名:
hook(module, prefix, keep_vars) -> None
注册的 hook 可用于在调用
state_dict
之前执行预处理。
- requires_grad_(requires_grad: bool = True) T ¶
改变是否 autograd 应记录此模块参数上的操作。
此方法就地设置参数的
requires_grad
属性。此方法对于微调时冻结模块部分或单独训练模型部分(例如 GAN 训练)很有帮助。
请参阅 局部禁用梯度计算,了解 `.requires_grad_()` 与其他可能与之混淆的类似机制的比较。
- 参数::
requires_grad (bool) – autograd 是否应记录此模块参数上的操作。默认值:
True
。- 返回值::
自身
- 返回值类型::
模块
- reset_out_keys()¶
将
out_keys
属性重置为其原始值。返回:具有其原始
out_keys
值的同一模块。示例
>>> from tensordict import TensorDict >>> from tensordict.nn import TensorDictModule, TensorDictSequential >>> import torch >>> mod = TensorDictModule(lambda x, y: (x+2, y+2), in_keys=["a", "b"], out_keys=["c", "d"]) >>> mod.select_out_keys("d") >>> td = TensorDict({"a": torch.zeros(()), "b": torch.ones(())}, []) >>> mod(td) TensorDict( fields={ a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), b: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), d: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False) >>> mod.reset_out_keys() >>> mod(td) TensorDict( fields={ a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), b: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), d: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False)
- reset_parameters_recursive(parameters: Optional[TensorDictBase] = None) Optional[TensorDictBase] ¶
递归地重置模块及其子模块的参数。
- 参数::
parameters (参数的 TensorDict,可选) – 如果设置为 None,模块将使用 self.parameters() 进行重置。否则,我们将就地重置 tensordict 中的参数。这对于参数未存储在模块本身中的函数式模块很有用。
- 返回值::
新参数的 tensordict,仅当 parameters 不为 None 时返回。
示例
>>> from tensordict.nn import TensorDictModule >>> from torch import nn >>> net = nn.Sequential(nn.Linear(2,3), nn.ReLU()) >>> old_param = net[0].weight.clone() >>> module = TensorDictModule(net, in_keys=['bork'], out_keys=['dork']) >>> module.reset_parameters() >>> (old_param == net[0].weight).any() tensor(False)
此方法也支持函数式参数采样
>>> from tensordict import TensorDict >>> from tensordict.nn import TensorDictModule >>> from torch import nn >>> net = nn.Sequential(nn.Linear(2,3), nn.ReLU()) >>> module = TensorDictModule(net, in_keys=['bork'], out_keys=['dork']) >>> params = TensorDict.from_module(module) >>> old_params = params.clone(recurse=True) >>> module.reset_parameters(params) >>> (old_params == params).any() False
- select_out_keys(*out_keys) TensorDictModuleBase ¶
选择在输出 tensordict 中找到的键。
这在想要摆脱复杂图中的中间键或当这些键的存在可能触发意外行为时很有用。
原始的
out_keys
仍可通过module.out_keys_source
访问。- 参数::
*out_keys (字符串序列 或 字符串元组) – 输出 tensordict 中应包含的 out_keys。
返回:同一模块,已就地修改并更新了
out_keys
。最简单的用法是与
TensorDictModule
一起使用。示例
>>> from tensordict import TensorDict >>> from tensordict.nn import TensorDictModule, TensorDictSequential >>> import torch >>> mod = TensorDictModule(lambda x, y: (x+2, y+2), in_keys=["a", "b"], out_keys=["c", "d"]) >>> td = TensorDict({"a": torch.zeros(()), "b": torch.ones(())}, []) >>> mod(td) TensorDict( fields={ a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), b: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), d: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False) >>> mod.select_out_keys("d") >>> td = TensorDict({"a": torch.zeros(()), "b": torch.ones(())}, []) >>> mod(td) TensorDict( fields={ a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), b: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), d: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False)
此功能也适用于分派参数: .. rubric:: 示例
>>> mod(torch.zeros(()), torch.ones(())) tensor(2.)
此更改将就地发生(即,返回的将是具有更新的 out_keys 列表的同一模块)。可以使用
TensorDictModuleBase.reset_out_keys()
方法将其还原。示例
>>> mod.reset_out_keys() >>> mod(TensorDict({"a": torch.zeros(()), "b": torch.ones(())}, [])) TensorDict( fields={ a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), b: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), d: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False)
这也适用于其他类,例如 Sequential: .. rubric:: 示例
>>> from tensordict.nn import TensorDictSequential >>> seq = TensorDictSequential( ... TensorDictModule(lambda x: x+1, in_keys=["x"], out_keys=["y"]), ... TensorDictModule(lambda x: x+1, in_keys=["y"], out_keys=["z"]), ... ) >>> td = TensorDict({"x": torch.zeros(())}, []) >>> seq(td) TensorDict( fields={ x: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), y: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), z: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False) >>> seq.select_out_keys("z") >>> td = TensorDict({"x": torch.zeros(())}, []) >>> seq(td) TensorDict( fields={ x: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), z: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False)
- set_extra_state(state: Any) None ¶
设置加载的 state_dict 中包含的额外状态。
此函数由
load_state_dict()
调用,用于处理在 state_dict 中找到的任何额外状态。如果需要在模块的 state_dict 中存储额外状态,请实现此函数以及相应的get_extra_state()
。- 参数::
state (dict) – 来自 state_dict 的额外状态。
- set_submodule(target: str, module: Module) None ¶
设置由
target
给定的子模块,如果存在则设置,否则抛出错误。例如,假设你有一个
nn.Module
A
,结构如下A( (net_b): Module( (net_c): Module( (conv): Conv2d(16, 33, kernel_size=(3, 3), stride=(2, 2)) ) (linear): Linear(in_features=100, out_features=200, bias=True) ) )
(图显示了一个
nn.Module
A
。A
有一个嵌套子模块net_b
,它本身有两个子模块net_c
和linear
。net_c
然后有一个子模块conv
。)要用新的子模块
Linear
覆盖Conv2d
,您可以调用set_submodule("net_b.net_c.conv", nn.Linear(33, 16))
。- 参数::
target – 要查找的子模块的完全限定字符串名称。(有关如何指定完全限定字符串,请参阅上述示例。)
module – 要设置的子模块。
- 抛出::
ValueError – 如果目标字符串为空。
AttributeError – 如果 target 字符串引用了无效路径或解析到的内容不是
nn.Module
- state_dict(*args, destination=None, prefix='', keep_vars=False)¶
返回一个字典,其中包含对模块整体状态的引用。
参数和持久缓冲区(例如,运行平均值)都包含在内。键对应于参数和缓冲区的名称。设置为
None
的参数和缓冲区不包含在内。注意
返回的对象是浅拷贝。它包含对模块参数和缓冲区的引用。
警告
目前,
state_dict()
也按顺序接受用于destination
、prefix
和keep_vars
的位置参数。然而,这正在被弃用,未来的版本中将强制使用关键字参数。警告
请避免使用参数
destination
,因为它不是为终端用户设计的。- 参数::
destination (dict, optional) – 如果提供,模块的状态将更新到此 dict 中并返回同一对象。否则,将创建一个
OrderedDict
并返回。默认值:None
。prefix (str, optional) – 添加到参数和缓冲区名称前缀以构成 state_dict 中的键。默认值:
''
。keep_vars (bool, optional) – 默认情况下,state dict 中返回的
Tensor
会从 autograd 中分离。如果设置为True
,则不会执行分离。默认值:False
。
- 返回值::
包含模块整体状态的字典。
- 返回值类型::
dict
示例
>>> # xdoctest: +SKIP("undefined vars") >>> module.state_dict().keys() ['bias', 'weight']
- to(*args, **kwargs)¶
移动和/或转换参数和缓冲区。
可以按以下方式调用:
- to(device=None, dtype=None, non_blocking=False)
- to(dtype, non_blocking=False)
- to(tensor, non_blocking=False)
- to(memory_format=torch.channels_last)
其签名类似于
torch.Tensor.to()
,但只接受浮点或复数dtype
。此外,此方法仅将浮点或复数参数和缓冲区转换为dtype
(如果给定)。整数参数和缓冲区将被移动到device
(如果给定),但其 dtype 不变。设置non_blocking
时,它会尝试尽可能异步地进行转换/移动,例如,将具有固定内存的 CPU Tensors 移动到 CUDA 设备。请参阅下面的示例。
注意
此方法会就地修改模块。
- 参数::
device (
torch.device
) – 此模块中参数和缓冲区的期望设备。dtype (
torch.dtype
) – 此模块中参数和缓冲区的期望浮点或复数 dtype。tensor (torch.Tensor) – 其 dtype 和 device 是此模块中所有参数和缓冲区的期望 dtype 和 device 的 Tensor。
memory_format (
torch.memory_format
) – 此模块中 4D 参数和缓冲区的期望内存格式(仅关键字参数)。
- 返回值::
自身
- 返回值类型::
模块
示例
>>> # xdoctest: +IGNORE_WANT("non-deterministic") >>> linear = nn.Linear(2, 2) >>> linear.weight Parameter containing: tensor([[ 0.1913, -0.3420], [-0.5113, -0.2325]]) >>> linear.to(torch.double) Linear(in_features=2, out_features=2, bias=True) >>> linear.weight Parameter containing: tensor([[ 0.1913, -0.3420], [-0.5113, -0.2325]], dtype=torch.float64) >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA1) >>> gpu1 = torch.device("cuda:1") >>> linear.to(gpu1, dtype=torch.half, non_blocking=True) Linear(in_features=2, out_features=2, bias=True) >>> linear.weight Parameter containing: tensor([[ 0.1914, -0.3420], [-0.5112, -0.2324]], dtype=torch.float16, device='cuda:1') >>> cpu = torch.device("cpu") >>> linear.to(cpu) Linear(in_features=2, out_features=2, bias=True) >>> linear.weight Parameter containing: tensor([[ 0.1914, -0.3420], [-0.5112, -0.2324]], dtype=torch.float16) >>> linear = nn.Linear(2, 2, bias=None).to(torch.cdouble) >>> linear.weight Parameter containing: tensor([[ 0.3741+0.j, 0.2382+0.j], [ 0.5593+0.j, -0.4443+0.j]], dtype=torch.complex128) >>> linear(torch.ones(3, 2, dtype=torch.cdouble)) tensor([[0.6122+0.j, 0.1150+0.j], [0.6122+0.j, 0.1150+0.j], [0.6122+0.j, 0.1150+0.j]], dtype=torch.complex128)
- to_empty(*, device: Optional[Union[int, str, device]], recurse: bool = True) T ¶
将参数和缓冲区移动到指定设备,但不复制存储。
- 参数::
device (
torch.device
) – 此模块中参数和缓冲区的期望设备。recurse (bool) – 子模块的参数和缓冲区是否应递归移动到指定设备。
- 返回值::
自身
- 返回值类型::
模块
- train(mode: bool = True) T ¶
将模块设置为训练模式。
这仅对某些模块有效。有关其在训练/评估模式下的行为(即,它们是否受到影响)的详细信息,请参阅特定模块的文档,例如
Dropout
、BatchNorm
等。- 参数::
mode (bool) – 是否设置为训练模式(
True
)或评估模式(False
)。默认值:True
。- 返回值::
自身
- 返回值类型::
模块
- type(dst_type: Union[dtype, str]) T ¶
将所有参数和缓冲区转换为
dst_type
。注意
此方法会就地修改模块。
- 参数::
dst_type (类型 或 字符串) – 期望的类型。
- 返回值::
自身
- 返回值类型::
模块
- xpu(device: Optional[Union[int, device]] = None) T ¶
将所有模型参数和缓冲区移动到 XPU。
这也会使相关参数和缓冲区成为不同的对象。因此,如果模块将在 XPU 上进行优化,则应在构建优化器之前调用此方法。
注意
此方法会就地修改模块。
- 参数::
device (整数,可选) – 如果指定,所有参数将复制到该设备
- 返回值::
自身
- 返回值类型::
模块
- zero_grad(set_to_none: bool = True) None ¶
重置所有模型参数的梯度。
请参阅
torch.optim.Optimizer
下的类似函数以获取更多背景信息。- 参数::
set_to_none (bool) – 将梯度设置为 None 而不是零。有关详细信息,请参阅
torch.optim.Optimizer.zero_grad()
。