RetrieveLogProb

class torchrl.envs.llm.transforms.RetrieveLogProb(model: LLMWrapperBase, *, log_probs_full_key: NestedKey | None = None, assistant_only: bool = True, tokenizer_kwargs: dict | None = None, detach: bool = True, device: torch.device | None = None, tokenizer: transformers.AutoTokenizer | None = None, padding_side: str = 'left')

A transform to retrieve log-probabilities from a model for KL divergence computation.

This transform computes log-probabilities from a reference model, which can then be used to compute KL divergence with another model’s log-probabilities. It’s designed to work with the RetrieveKL and KLComputation transforms.

参数:

model (LLMWrapperBase) – the model to use to compute the log-probs.

关键字参数:

• log_probs_full_key (NestedKey) – the key where the log-probs are stored. If not provided, the key will be retrieved from the model’s log_probs_key attribute (i.e., (model.log_probs_key, “full”)).

• assistant_only (bool) –

  whether to zero out the log-probs of the non-assistant tokens (i.e., steps of history where the role is not “assistant”). Defaults to True.

  注意

  When assistant_only=True, the model must have input_mode=’history’ to properly identify assistant tokens. For other input modes (“text” or “tokens”), set assistant_only=False. This ensures users are conscious of the limitation that assistant token identification requires structured conversation history.

• tokenizer_kwargs (dict) – the keyword arguments to pass to the tokenizer to be used to apply the chat template to the history when assistant_only is True. To control the tokenization in the ref_model, pass the tokenizer kwargs to the ref_model constructor. Defaults to {“return_assistant_tokens_mask”: True, “tokenize”: True, “return_dict”: True, “padding”: False, “add_generation_prompt”: False}.

• tokenizer (transformers.AutoTokenizer) – the tokenizer to be used to tokenize the input and compute the assitant mask. If not provided, the tokenizer will be inferred from the ref_model.

• detach (bool) – whether to exclude the log-probs from the gradient computation. Defaults to True.

• device (torch.device) – the device to use for tensor creation. Defaults to None.

• padding_side (str) – 使用 pad_sequence 时填充的侧边。默认为 “left”。

示例

>>> from torchrl.data.llm import History
>>> from torchrl.modules.llm import TransformersWrapper
>>> from torchrl.modules.llm.policies import ChatHistory
>>> from transformers import AutoTokenizer, OPTConfig, OPTForCausalLM
>>> from tensordict import TensorDict, set_list_to_stack
>>> import torch
>>>
>>> # Set up list to stack for History
>>> set_list_to_stack(True).set()
>>>
>>> # Create chat data
>>> chats = [
...     [
...         {"role": "system", "content": "You are a helpful assistant."},
...         {"role": "user", "content": "Hello, how are you?"},
...         {"role": "assistant", "content": "I'm doing well, thank you!"},
...     ],
...     [
...         {"role": "system", "content": "You are a helpful assistant."},
...         {"role": "user", "content": "What's the weather like?"},
...         {"role": "assistant", "content": "I can't check the weather for you."},
...     ],
... ]
>>> history = History.from_chats(chats)
>>> print(f"Created history with shape: {history.shape}")
Created history with shape: torch.Size([2, 3])
>>>
>>> # Setup tokenizer and model
>>> tokenizer = AutoTokenizer.from_pretrained("facebook/opt-125m")
>>> tokenizer.pad_token = tokenizer.eos_token
>>> model = OPTForCausalLM(OPTConfig()).eval()
>>>
>>> # Create reference model
>>> ref_model = TransformersWrapper(
...     model,
...     tokenizer=tokenizer,
...     input_mode="history",
...     generate=False,
...     return_log_probs=True,
...     pad_output=True,
... )
>>>
>>> # Create the RetrieveLogProb transform
>>> transform = RetrieveLogProb(
...     ref_model,
...     assistant_only=True,
...     tokenizer=tokenizer,
... )
>>>
>>> # Prepare data using ChatHistory
>>> chat_history = ChatHistory(full=history)
>>> data = TensorDict(history=chat_history, batch_size=(2,))
>>>
>>> # Apply the transform to get reference log probabilities
>>> result = transform(data)
>>> log_probs_key = (ref_model.log_probs_key, "full")
>>> ref_log_probs = result.get(log_probs_key)
>>> print(f"Log-probs shape: {ref_log_probs.shape}")
Log-probs shape: torch.Size([2, 26])

注意

By default, the log-probabilities are stored as a list of tensors (one per sample, with variable length). Use as_padded_tensor=True in .get() to obtain a batchable tensor (with padding). The reference log probabilities are computed only for assistant tokens when assistant_only=True.

Input Mode Compatibility: - When assistant_only=True (default), the model must have input_mode=’history’ to properly identify assistant tokens. - When assistant_only=False, the transform works with any input mode (“history”, “text”, or “tokens”). - This design ensures users are conscious of the limitation that assistant token identification requires structured conversation history.

另请参阅

RetrieveKL: A higher-level transform that combines two RetrieveLogProb instances with KLComputation. KLComputation: A transform that computes KL divergence between two log-prob tensors. KLRewardTransform: A legacy transform for KL reward computation (use RetrieveKL instead).

add_module(name: str, module: Optional[Module]) → None

将子模块添加到当前模块。

可以使用给定的名称作为属性访问该模块。

参数:

• name (str) – 子模块的名称。子模块可以通过给定名称从此模块访问

• module (Module) – 要添加到模块中的子模块。

apply(fn: Callable[[Module], None]) → T

将 fn 递归应用于每个子模块（由 .children() 返回）以及自身。

典型用法包括初始化模型的参数（另请参阅 torch.nn.init）。

参数:

fn (Module -> None) – 要应用于每个子模块的函数

返回:

self

返回类型:

模块

示例

>>> @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 数据类型。

注意

此方法就地修改模块。

返回:

self

返回类型:

模块

buffers(recurse: bool = True) → Iterator[Tensor]

返回模块缓冲区的迭代器。

参数:

recurse (bool) – 如果为 True，则会产生此模块及其所有子模块的 buffer。否则，仅会产生此模块的直接成员 buffer。

产生:

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)

children() → Iterator[Module]

返回直接子模块的迭代器。

产生:

Module – 子模块

close()

关闭转换。

property collector: DataCollectorBase | None

返回与容器关联的收集器（如果存在）。

每当变换需要了解收集器或与之关联的策略时，都可以使用此属性。请确保仅在未嵌套在子进程中的变换上调用此属性。收集器引用不会传递给 ParallelEnv 或类似的批处理环境的 worker。

Make sure to call this property only on transforms that are not nested in sub-processes. The collector reference will not be passed to the workers of a ParallelEnv or similar batched environments.

compile(*args, **kwargs)

Compile this Module’s forward using torch.compile().

This Module’s __call__ method is compiled and all arguments are passed as-is to torch.compile().

See torch.compile() for details on the arguments for this function.

property container: EnvBase | None

返回包含该变换的环境。

示例

>>> from torchrl.envs import TransformedEnv, Compose, RewardSum, StepCounter
>>> from torchrl.envs.libs.gym import GymEnv
>>> env = TransformedEnv(GymEnv("Pendulum-v1"), Compose(RewardSum(), StepCounter()))
>>> env.transform[0].container is env
True

cpu() → T

将所有模型参数和缓冲区移动到 CPU。

注意

此方法就地修改模块。

返回:

self

返回类型:

模块

cuda(device: Optional[Union[int, device]] = None) → T

将所有模型参数和缓冲区移动到 GPU。

这也会使相关的参数和缓冲区成为不同的对象。因此，如果模块在优化时将驻留在 GPU 上，则应在构建优化器之前调用此函数。

注意

此方法就地修改模块。

参数:

device (int, optional) – 如果指定，所有参数将复制到该设备

返回:

self

返回类型:

模块

double() → T

将所有浮点参数和缓冲区转换为 double 数据类型。

注意

此方法就地修改模块。

返回:

self

返回类型:

模块

eval() → T

将模块设置为评估模式。

This has an effect only on certain modules. See the documentation of particular modules for details of their behaviors in training/evaluation mode, i.e. whether they are affected, e.g. Dropout, BatchNorm, etc.

This is equivalent with self.train(False).

See Locally disabling gradient computation for a comparison between .eval() and several similar mechanisms that may be confused with it.

返回:

self

返回类型:

模块

extra_repr() → str

返回模块的额外表示。

要打印自定义额外信息，您应该在自己的模块中重新实现此方法。单行和多行字符串均可接受。

float() → T

将所有浮点参数和缓冲区转换为 float 数据类型。

注意

此方法就地修改模块。

返回:

self

返回类型:

模块

forward(tensordict: TensorDictBase) → TensorDictBase

读取输入 tensordict，并对选定的键应用转换。

默认情况下，此方法

• 直接调用 _apply_transform()。

• 不调用 _step() 或 _call()。

This method is not called within env.step at any point. However, is is called within sample().

注意

forward also works with regular keyword arguments using dispatch to cast the args names to the keys.

示例

>>> class TransformThatMeasuresBytes(Transform):
...     '''Measures the number of bytes in the tensordict, and writes it under `"bytes"`.'''
...     def __init__(self):
...         super().__init__(in_keys=[], out_keys=["bytes"])
...
...     def forward(self, tensordict: TensorDictBase) -> TensorDictBase:
...         bytes_in_td = tensordict.bytes()
...         tensordict["bytes"] = bytes
...         return tensordict
>>> t = TransformThatMeasuresBytes()
>>> env = env.append_transform(t) # works within envs
>>> t(TensorDict(a=0))  # Works offline too.

get_buffer(target: str) → Tensor

返回由 target 给定的缓冲区（如果存在），否则抛出错误。

See the docstring for get_submodule for a more detailed explanation of this method’s functionality as well as how to correctly specify target.

参数:

target – 要查找的 buffer 的完全限定字符串名称。（要指定完全限定字符串，请参阅 get_submodule。）

返回:

由 target 引用的缓冲区

返回类型:

torch.Tensor

抛出:

AttributeError – 如果目标字符串引用了无效路径或解析为非 buffer 对象。

get_extra_state() → Any

返回要包含在模块 state_dict 中的任何额外状态。

Implement this and a corresponding set_extra_state() for your module if you need to store extra state. This function is called when building the module’s state_dict().

Note that extra state should be picklable to ensure working serialization of the state_dict. We only provide backwards compatibility guarantees for serializing Tensors; other objects may break backwards compatibility if their serialized pickled form changes.

返回:

要存储在模块 state_dict 中的任何额外状态

返回类型:

对象

get_parameter(target: str) → Parameter

如果存在，返回由 target 给定的参数，否则抛出错误。

See the docstring for get_submodule for a more detailed explanation of this method’s functionality as well as how to correctly specify target.

参数:

target – 要查找的 Parameter 的完全限定字符串名称。（要指定完全限定字符串，请参阅 get_submodule。）

返回:

由 target 引用的参数

返回类型:

torch.nn.Parameter

抛出:

AttributeError – 如果目标字符串引用了无效路径或解析为非 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)
    )
)

(The diagram shows an nn.Module A. A which has a nested submodule net_b, which itself has two submodules net_c and linear. net_c then has a submodule conv.)

To check whether or not we have the linear submodule, we would call get_submodule("net_b.linear"). To check whether we have the conv submodule, we would call get_submodule("net_b.net_c.conv").

The runtime of get_submodule is bounded by the degree of module nesting in target. A query against named_modules achieves the same result, but it is O(N) in the number of transitive modules. So, for a simple check to see if some submodule exists, get_submodule should always be used.

参数:

target – 要查找的子模块的完全限定字符串名称。（要指定完全限定字符串，请参阅上面的示例。）

返回:

由 target 引用的子模块

返回类型:

torch.nn.Module

抛出:

AttributeError – If at any point along the path resulting from the target string the (sub)path resolves to a non-existent attribute name or an object that is not an instance of nn.Module.

half() → T

将所有浮点参数和缓冲区转换为 half 数据类型。

注意

此方法就地修改模块。

返回:

self

返回类型:

模块

init(tensordict) → None

运行转换的初始化步骤。

inv(tensordict: TensorDictBase = None) → TensorDictBase

读取输入 tensordict，并对选定的键应用逆变换。

默认情况下，此方法

• 直接调用 _inv_apply_transform()。

• 不调用 _inv_call()。

注意

inv also works with regular keyword arguments using dispatch to cast the args names to the keys.

注意

inv is called by extend().

ipu(device: Optional[Union[int, device]] = None) → T

将所有模型参数和缓冲区移动到 IPU。

这也会使关联的参数和缓冲区成为不同的对象。因此，如果模块在优化时将驻留在 IPU 上，则应在构建优化器之前调用它。

注意

此方法就地修改模块。

参数:

device (int, optional) – 如果指定，所有参数将复制到该设备

返回:

self

返回类型:

模块

load_state_dict(state_dict: Mapping[str, Any], strict: bool = True, assign: bool = False)

Copy parameters and buffers from state_dict into this module and its descendants.

If strict is True, then the keys of state_dict must exactly match the keys returned by this module’s state_dict() function.

警告

If assign is True the optimizer must be created after the call to load_state_dict unless get_swap_module_params_on_conversion() is True.

参数:

• state_dict (dict) – 包含参数和持久 buffer 的字典。

• strict (bool, optional) – whether to strictly enforce that the keys in state_dict match the keys returned by this module’s state_dict() function. Default: True

• assign (bool, optional) – When set to False, the properties of the tensors in the current module are preserved whereas setting it to True preserves properties of the Tensors in the state dict. The only exception is the requires_grad field of Default: ``False`

返回:

• missing_keys 是一个包含任何预期键的 str 列表。

  在提供的 state_dict 中缺失的任何键的字符串列表。

• unexpected_keys 是一个包含不匹配的键的 str 列表。

  不期望但在提供的 state_dict 中存在的键。

返回类型:

NamedTuple with missing_keys and unexpected_keys fields

注意

If a parameter or buffer is registered as None and its corresponding key exists in state_dict, load_state_dict() will raise a 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 (int, optional) – 如果指定，所有参数将复制到该设备

返回:

self

返回类型:

模块

named_buffers(prefix: str = '', recurse: bool = True, remove_duplicate: bool = True) → Iterator[tuple[str, torch.Tensor]

返回模块缓冲区上的迭代器，同时生成缓冲区的名称和缓冲区本身。

参数:

• prefix (str) – 为所有 buffer 名称添加前缀。

• recurse (bool, optional) – 如果为 True，则会生成此模块及其所有子模块的 buffers。否则，仅生成此模块直接成员的 buffers。默认为 True。

• remove_duplicate (bool, optional) – 是否在结果中删除重复的 buffers。默认为 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 – 用于存储已添加到结果中的模块集合的 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, torch.nn.parameter.Parameter]

返回模块参数的迭代器，同时生成参数的名称和参数本身。

参数:

• prefix (str) – 为所有参数名称添加前缀。

• recurse (bool) – 如果为 True，则会生成此模块及其所有子模块的参数。否则，仅生成此模块直接成员的参数。

• remove_duplicate (bool, optional) – 是否在结果中删除重复的参数。默认为 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)

property parent: TransformedEnv | None

返回变换的父环境。

父环境是包含直到当前变换的所有变换的环境。

示例

>>> from torchrl.envs import TransformedEnv, Compose, RewardSum, StepCounter
>>> from torchrl.envs.libs.gym import GymEnv
>>> env = TransformedEnv(GymEnv("Pendulum-v1"), Compose(RewardSum(), StepCounter()))
>>> env.transform[1].parent
TransformedEnv(
    env=GymEnv(env=Pendulum-v1, batch_size=torch.Size([]), device=cpu),
    transform=Compose(
            RewardSum(keys=['reward'])))

register_backward_hook(hook: Callable[[Module, Union[tuple[torch.Tensor, ...], Tensor], Union[tuple[torch.Tensor, ...], Tensor]]) → RemovableHandle

在模块上注册一个反向传播钩子。

This function is deprecated in favor of register_full_backward_hook() and the behavior of this function will change in future versions.

返回:

一个句柄，可用于通过调用 handle.remove() 来移除添加的钩子

返回类型:

torch.utils.hooks.RemovableHandle

register_buffer(name: str, tensor: Optional[Tensor], persistent: bool = True) → None

向模块添加一个缓冲区。

This is typically used to register a buffer that should not to be considered a model parameter. For example, BatchNorm’s running_mean is not a parameter, but is part of the module’s state. Buffers, by default, are persistent and will be saved alongside parameters. This behavior can be changed by setting persistent to False. The only difference between a persistent buffer and a non-persistent buffer is that the latter will not be a part of this module’s state_dict.

可以使用给定名称作为属性访问缓冲区。

参数:

• name (str) – buffer 的名称。可以使用给定的名称从此模块访问 buffer

• tensor (Tensor or None) – buffer to be registered. If None, then operations that run on buffers, such as cuda, are ignored. If None, the buffer is not included in the module’s state_dict.

• persistent (bool) – whether the buffer is part of this module’s 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

在模块上注册一个前向钩子。

The hook will be called every time after forward() has computed an output.

If with_kwargs is False or not specified, the input contains only the positional arguments given to the module. Keyword arguments won’t be passed to the hooks and only to the forward. The hook can modify the output. It can modify the input inplace but it will not have effect on forward since this is called after forward() is called. The hook should have the following signature

hook(module, args, output) -> None or modified output

If with_kwargs is True, the forward hook will be passed the kwargs given to the forward function and be expected to return the output possibly modified. The hook should have the following signature

hook(module, args, kwargs, output) -> None or modified output

参数:

• hook (Callable) – 用户定义的待注册钩子。

• prepend (bool) – If True, the provided hook will be fired before all existing forward hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing forward hooks on this torch.nn.Module. Note that global forward hooks registered with register_module_forward_hook() will fire before all hooks registered by this method. Default: False

• with_kwargs (bool) – If True, the hook will be passed the kwargs given to the forward function. Default: False

• always_call (bool) – If True the hook will be run regardless of whether an exception is raised while calling the Module. Default: 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

在模块上注册一个前向预钩子。

The hook will be called every time before forward() is invoked.

If with_kwargs is false or not specified, the input contains only the positional arguments given to the module. Keyword arguments won’t be passed to the hooks and only to the forward. The hook can modify the input. User can either return a tuple or a single modified value in the hook. We will wrap the value into a tuple if a single value is returned (unless that value is already a tuple). The hook should have the following signature

hook(module, args) -> None or modified input

If with_kwargs is true, the forward pre-hook will be passed the kwargs given to the forward function. And if the hook modifies the input, both the args and kwargs should be returned. The hook should have the following signature

hook(module, args, kwargs) -> None or a tuple of modified input and kwargs

参数:

• hook (Callable) – 用户定义的待注册钩子。

• prepend (bool) – If true, the provided hook will be fired before all existing forward_pre hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing forward_pre hooks on this torch.nn.Module. Note that global forward_pre hooks registered with register_module_forward_pre_hook() will fire before all hooks registered by this method. Default: False

• with_kwargs (bool) – If true, the hook will be passed the kwargs given to the forward function. Default: False

返回:

一个句柄，可用于通过调用 handle.remove() 来移除添加的钩子

返回类型:

torch.utils.hooks.RemovableHandle

register_full_backward_hook(hook: Callable[[Module, Union[tuple[torch.Tensor, ...], Tensor], Union[tuple[torch.Tensor, ...], Tensor]], Union[None, tuple[torch.Tensor, ...], Tensor]], prepend: bool = False) → RemovableHandle

在模块上注册一个反向传播钩子。

hook 将在计算模块的梯度时被调用，即只有在计算模块输出的梯度时 hook 才会执行。hook 应具有以下签名

hook(module, grad_input, grad_output) -> tuple(Tensor) or None

The grad_input and grad_output are tuples that contain the gradients with respect to the inputs and outputs respectively. The hook should not modify its arguments, but it can optionally return a new gradient with respect to the input that will be used in place of grad_input in subsequent computations. grad_input will only correspond to the inputs given as positional arguments and all kwarg arguments are ignored. Entries in grad_input and grad_output will be None for all non-Tensor arguments.

由于技术原因，当此钩子应用于模块时，其前向函数将接收传递给模块的每个张量的视图。类似地，调用者将接收模块前向函数返回的每个张量的视图。

警告

使用反向传播钩子时不允许就地修改输入或输出，否则将引发错误。

参数:

• hook (Callable) – 要注册的用户定义钩子。

• prepend (bool) – If true, the provided hook will be fired before all existing backward hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing backward hooks on this torch.nn.Module. Note that global backward hooks registered with register_module_full_backward_hook() will fire before all hooks registered by this method.

返回:

一个句柄，可用于通过调用 handle.remove() 来移除添加的钩子

返回类型:

torch.utils.hooks.RemovableHandle

register_full_backward_pre_hook(hook: Callable[[Module, Union[tuple[torch.Tensor, ...], Tensor]], Union[None, tuple[torch.Tensor, ...], Tensor]], prepend: bool = False) → RemovableHandle

在模块上注册一个反向预钩子。

每次计算模块的梯度时，将调用此钩子。钩子应具有以下签名

hook(module, grad_output) -> tuple[Tensor] or None

The grad_output is a tuple. The hook should not modify its arguments, but it can optionally return a new gradient with respect to the output that will be used in place of grad_output in subsequent computations. Entries in grad_output will be None for all non-Tensor arguments.

由于技术原因，当此钩子应用于模块时，其前向函数将接收传递给模块的每个张量的视图。类似地，调用者将接收模块前向函数返回的每个张量的视图。

警告

使用反向传播钩子时不允许就地修改输入，否则将引发错误。

参数:

• hook (Callable) – 要注册的用户定义钩子。

• prepend (bool) – If true, the provided hook will be fired before all existing backward_pre hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing backward_pre hooks on this torch.nn.Module. Note that global backward_pre hooks registered with register_module_full_backward_pre_hook() will fire before all hooks registered by this method.

返回:

一个句柄，可用于通过调用 handle.remove() 来移除添加的钩子

返回类型:

torch.utils.hooks.RemovableHandle

register_load_state_dict_post_hook(hook)

注册一个后钩子，用于在模块的 load_state_dict() 被调用后运行。

它应该具有以下签名：

hook(module, incompatible_keys) -> None

The module argument is the current module that this hook is registered on, and the incompatible_keys argument is a NamedTuple consisting of attributes missing_keys and unexpected_keys. missing_keys is a list of str containing the missing keys and unexpected_keys is a list of str containing the unexpected keys.

如果需要，可以就地修改给定的 incompatible_keys。

Note that the checks performed when calling load_state_dict() with strict=True are affected by modifications the hook makes to missing_keys or unexpected_keys, as expected. Additions to either set of keys will result in an error being thrown when strict=True, and clearing out both missing and unexpected keys will avoid an error.

返回:

一个句柄，可用于通过调用 handle.remove() 来移除添加的钩子

返回类型:

torch.utils.hooks.RemovableHandle

register_load_state_dict_pre_hook(hook)

注册一个预钩子，用于在模块的 load_state_dict() 被调用之前运行。

它应该具有以下签名：

hook(module, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs) -> None # noqa: B950

参数:

hook (Callable) – 在加载状态字典之前将调用的可调用钩子。

register_module(name: str, module: Optional[Module]) → None

Alias for 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(module, state_dict, prefix, local_metadata) -> None

注册的钩子可以就地修改 state_dict。

register_state_dict_pre_hook(hook)

为 state_dict() 方法注册一个前钩子。

它应该具有以下签名：

hook(module, prefix, keep_vars) -> None

注册的钩子可用于在进行 state_dict 调用之前执行预处理。

requires_grad_(requires_grad: bool = True) → T

更改自动梯度是否应记录此模块中参数的操作。

此方法就地设置参数的 requires_grad 属性。

此方法有助于冻结模块的一部分以进行微调或单独训练模型的一部分（例如，GAN 训练）。

请参阅 局部禁用梯度计算 以比较 .requires_grad_() 和几个可能与之混淆的类似机制。

参数:

requires_grad (bool) – 自动求导是否应记录此模块上的参数操作。默认为 True。

返回:

self

返回类型:

模块

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, strict: bool = False) → None

如果存在，设置由 target 给定的子模块，否则抛出错误。

注意

如果 strict 设置为 False（默认值），则该方法将替换现有子模块或在父模块存在时创建新子模块。如果 strict 设置为 True，则该方法只会尝试替换现有子模块，如果子模块不存在则会抛出错误。

例如，假设您有一个 nn.Module A，它看起来像这样

A(
    (net_b): Module(
        (net_c): Module(
            (conv): Conv2d(3, 3, 3)
        )
        (linear): Linear(3, 3)
    )
)

(图示显示了一个 nn.Module A。 A 有一个嵌套子模块 net_b，它本身有两个子模块 net_c 和 linear。 net_c 然后有一个子模块 conv。)

要用新的 Linear 子模块覆盖 Conv2d，您可以调用 set_submodule("net_b.net_c.conv", nn.Linear(1, 1))，其中 strict 可以是 True 或 False。

要向现有 net_b 模块添加新的 Conv2d 子模块，您需要调用 set_submodule("net_b.conv", nn.Conv2d(1, 1, 1))。

在上面的示例中，如果您设置 strict=True 并调用 set_submodule("net_b.conv", nn.Conv2d(1, 1, 1), strict=True)，则会引发 AttributeError，因为 net_b 中不存在名为 conv 的子模块。

参数:

• target – 要查找的子模块的完全限定字符串名称。（要指定完全限定字符串，请参阅上面的示例。）

• module – 要设置子模块的对象。

• strict – 如果为 False，则该方法将替换现有子模块或在父模块存在时创建新子模块。如果为 True，则该方法只会尝试替换现有子模块，如果子模块尚不存在则会抛出错误。

抛出:

• ValueError – 如果 target 字符串为空，或者 module 不是 nn.Module 的实例。

• AttributeError – 如果沿 target 字符串产生的路径中的任何一点，(子)路径解析为不存在的属性名，或者一个不是 nn.Module 实例的对象。

share_memory() → T

请参阅 torch.Tensor.share_memory_()。

state_dict(*args, destination=None, prefix='', keep_vars=False)

返回一个字典，其中包含对模块整个状态的引用。

参数和持久缓冲区（例如，运行平均值）都包含在内。键是相应的参数和缓冲区名称。设置为 None 的参数和缓冲区不包含在内。

注意

返回的对象是浅拷贝。它包含对模块参数和缓冲区的引用。

警告

当前 state_dict() 还接受 destination、prefix 和 keep_vars 的位置参数。但是，这正在被弃用，并且在将来的版本中将强制使用关键字参数。

警告

请避免使用参数 destination，因为它不是为最终用户设计的。

参数:

• destination (dict, optional) – 如果提供，模块的状态将被更新到字典中，并返回相同的对象。否则，将创建一个 OrderedDict 并返回。默认值：None。

• prefix (str, optional) – a prefix added to parameter and buffer names to compose the keys in state_dict. Default: ''。

• keep_vars (bool, optional) – 默认情况下，Tensor s 返回在 state_dict 中是与 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 Tensor 移动到 CUDA 设备。

有关示例，请参阅下文。

注意

此方法就地修改模块。

参数:

• device (torch.device) – 此模块中参数和缓冲区的目标设备

• dtype (torch.dtype) – 此模块中参数和缓冲区的目标浮点数或复数 dtype

• tensor (torch.Tensor) – 其 dtype 和设备是此模块中所有参数和缓冲区的目标 dtype 和设备

• memory_format (torch.memory_format) – 4D 参数和缓冲区的目标内存格式（仅关键字参数）

返回:

self

返回类型:

模块

示例

>>> # 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) – 是否递归地将子模块的参数和缓冲区移动到指定设备。

返回:

self

返回类型:

模块

train(mode: bool = True) → T

将模块设置为训练模式。

这仅对某些模块有效。有关它们在训练/评估模式下的行为的详细信息，请参阅特定模块的文档，例如它们是否受影响，例如 Dropout、BatchNorm 等。

参数:

mode (bool) – 设置训练模式（True）还是评估模式（False）。默认值：True。

返回:

self

返回类型:

模块

transform_action_spec(action_spec: TensorSpec) → TensorSpec

转换动作规范，使结果规范与变换映射匹配。

参数:

action_spec (TensorSpec) – 变换前的规范

返回:

转换后的预期规范

transform_done_spec(done_spec: TensorSpec) → TensorSpec

变换 done spec，使结果 spec 与变换映射匹配。

参数:

done_spec (TensorSpec) – 变换前的 spec

返回:

转换后的预期规范

transform_env_batch_size(batch_size: Size)

转换父环境的 batch-size。

transform_env_device(device: device)

转换父环境的 device。

transform_input_spec(input_spec: TensorSpec) → TensorSpec

转换输入规范，使结果规范与转换映射匹配。

参数:

input_spec (TensorSpec) – 转换前的规范

返回:

转换后的预期规范

transform_observation_spec(observation_spec: Composite) → Composite

转换观察规范，使结果规范与转换映射匹配。

参数:

observation_spec (TensorSpec) – 转换前的规范

返回:

转换后的预期规范

transform_output_spec(output_spec: Composite) → Composite

转换输出规范，使结果规范与转换映射匹配。

此方法通常应保持不变。更改应通过 transform_observation_spec()、transform_reward_spec() 和 transform_full_done_spec() 来实现。 :param output_spec: 转换前的 spec :type output_spec: TensorSpec

返回:

转换后的预期规范

transform_reward_spec(reward_spec: TensorSpec) → TensorSpec

转换奖励的 spec，使其与变换映射匹配。

参数:

reward_spec (TensorSpec) – 变换前的 spec

返回:

转换后的预期规范

transform_state_spec(state_spec: TensorSpec) → TensorSpec

转换状态规范，使结果规范与变换映射匹配。

参数:

state_spec (TensorSpec) – 变换前的规范

返回:

转换后的预期规范

type(dst_type: Union[dtype, str]) → T

将所有参数和缓冲区转换为 dst_type。

注意

此方法就地修改模块。

参数:

dst_type (type or string) – 目标类型

返回:

self

返回类型:

模块

xpu(device: Optional[Union[int, device]] = None) → T

将所有模型参数和缓冲区移动到 XPU。

这也会使关联的参数和缓冲区成为不同的对象。因此，如果模块在优化时将驻留在 XPU 上，则应在构建优化器之前调用它。

注意

此方法就地修改模块。

参数:

device (int, optional) – 如果指定，所有参数将复制到该设备

返回:

self

返回类型:

模块

zero_grad(set_to_none: bool = True) → None

重置所有模型参数的梯度。

请参阅 torch.optim.Optimizer 下的类似函数以获取更多上下文。

参数:

set_to_none (bool) – 不要设置为零，而是将梯度设置为 None。详情请参阅 torch.optim.Optimizer.zero_grad()。