tensordict.nn.distributions.NormalParamExtractor

類 tensordict.nn.distributions.NormalParamExtractor(scale_mapping: str = 'biased_softplus_1.0', scale_lb: Number = 0.0001)

一個非引數化的 nn.Module，它將其輸入拆分為 loc 和 scale 引數。

使用指定的 scale_mapping 將 scale 引數對映到正值。

引數：

• scale_mapping (str, 可選) – 用於標準差的正向對映函式。預設值 = "biased_softplus_1.0" (即帶有偏置的 softplus 對映，使得 fn(0.0) = 1.0) 選項："softplus"、"exp"、"relu"、"biased_softplus_1" 或 "none" (無對映)。有關更多詳細資訊，請參閱 mappings()。

• scale_lb (Number, 可選) – 方差可以取的最小值。預設值為 1e-4。

示例

>>> import torch
>>> from tensordict.nn.distributions import NormalParamExtractor
>>> from torch import nn
>>> module = nn.Linear(3, 4)
>>> normal_params = NormalParamExtractor()
>>> tensor = torch.randn(3)
>>> loc, scale = normal_params(module(tensor))
>>> print(loc.shape, scale.shape)
torch.Size([2]) torch.Size([2])
>>> assert (scale > 0).all()
>>> # with modules that return more than one tensor
>>> module = nn.LSTM(3, 4)
>>> tensor = torch.randn(4, 2, 3)
>>> loc, scale, others = normal_params(*module(tensor))
>>> print(loc.shape, scale.shape)
torch.Size([4, 2, 2]) torch.Size([4, 2, 2])
>>> assert (scale > 0).all()