LayerNorm¶
- class nvtripy.LayerNorm(normalized_shape: int | Sequence[int], dtype: dtype = float32, eps: float = 1e-05)[source]¶
Bases:
Module
Applies layer normalization over the input tensor:
\(\text{LayerNorm}(x) = \Large \frac{x - \bar{x}}{ \sqrt{\sigma^2 + \epsilon}} \normalsize * \gamma + \beta\)
where \(\bar{x}\) is the mean and \(\sigma^2\) is the variance.
The mean and standard deviation are calculated over the last \(D\) dimensions, where \(D\) is the dimension of \(\text{normalized_shape}\).
- Parameters:
normalized_shape (Sequence[int]) – The size of the feature dimension of the input over which normalization is performed. If a single integer is provided, it will be unsqueezed to a 1 dimensional shape.
dtype (dtype) – The data type to use for the weight and bias parameters.
eps (float) – \(\epsilon\) value to prevent division by zero.
Example
1layer_norm = tp.LayerNorm(3) 2 3layer_norm.weight = tp.iota(layer_norm.weight.shape) 4layer_norm.bias = tp.iota(layer_norm.bias.shape) 5 6input = tp.iota((2, 3), dim=1) 7output = layer_norm(input)
Local Variables¶>>> layer_norm LayerNorm( weight: Parameter = (shape=(3,), dtype=float32), bias: Parameter = (shape=(3,), dtype=float32), ) >>> layer_norm.state_dict() { weight: tensor([0, 1, 2], dtype=float32, loc=gpu:0, shape=(3,)), bias: tensor([0, 1, 2], dtype=float32, loc=gpu:0, shape=(3,)), } >>> input tensor( [[0, 1, 2], [0, 1, 2]], dtype=float32, loc=gpu:0, shape=(2, 3)) >>> output tensor( [[0, 1, 4.44947], [0, 1, 4.44947]], dtype=float32, loc=gpu:0, shape=(2, 3))
- __call__(*args: Any, **kwargs: Any) Any ¶
Calls the module with the specified arguments.
- Parameters:
*args (Any) – Positional arguments to the module.
**kwargs (Any) – Keyword arguments to the module.
- Returns:
The outputs computed by the module.
- Return type:
Any
Example
1class Module(tp.Module): 2 def forward(self, x): 3 return tp.relu(x) 4 5 6module = Module() 7 8input = tp.arange(-3, 3) 9out = module(input) # Note that we do not call `forward` directly.
Local Variables¶>>> module Module( ) >>> module.state_dict() {} >>> input tensor([-3, -2, -1, 0, 1, 2], dtype=float32, loc=gpu:0, shape=(6,)) >>> out tensor([0, 0, 0, 0, 1, 2], dtype=float32, loc=gpu:0, shape=(6,))
- load_state_dict(state_dict: Dict[str, Tensor], strict: bool = True) Tuple[Set[str], Set[str]] ¶
Loads parameters from the provided
state_dict
into the current module. This will recurse over any nested child modules.- Parameters:
- Returns:
missing_keys: keys that are expected by this module but not provided in
state_dict
.unexpected_keys: keys that are not expected by this module but provided in
state_dict
.
- Return type:
A
tuple
of twoset
s of strings representing
Example
1class MyModule(tp.Module): 2 def __init__(self): 3 super().__init__() 4 self.param = tp.ones((2,), dtype=tp.float32) 5 6 7module = MyModule() 8 9print(f"Before: {module.param}") 10 11module.load_state_dict({"param": tp.zeros((2,), dtype=tp.float32)}) 12 13print(f"After: {module.param}")
Output¶Before: tensor([1, 1], dtype=float32, loc=gpu:0, shape=(2,)) After: tensor([0, 0], dtype=float32, loc=gpu:0, shape=(2,))
See also
- named_children() Iterator[Tuple[str, Module]] ¶
Returns an iterator over immediate children of this module, yielding tuples containing the name of the child module and the child module itself.
- Returns:
An iterator over tuples containing the name of the child module and the child module itself.
- Return type:
Iterator[Tuple[str, Module]]
Example
1class StackedLinear(tp.Module): 2 def __init__(self): 3 super().__init__() 4 self.linear1 = tp.Linear(2, 2) 5 self.linear2 = tp.Linear(2, 2) 6 7 8stacked_linear = StackedLinear() 9 10for name, module in stacked_linear.named_children(): 11 print(f"{name}: {type(module).__name__}")
Output¶linear1: Linear linear2: Linear
- named_parameters() Iterator[Tuple[str, Tensor]] ¶
- Returns:
An iterator over tuples containing the name of a parameter and the parameter itself.
- Return type:
Iterator[Tuple[str, Tensor]]
Example
1class MyModule(tp.Module): 2 def __init__(self): 3 super().__init__() 4 self.alpha = tp.Tensor(1) 5 self.beta = tp.Tensor(2) 6 7 8linear = MyModule() 9 10for name, parameter in linear.named_parameters(): 11 print(f"{name}: {parameter}")
Output¶alpha: tensor(1, dtype=int32, loc=cpu:0, shape=()) beta: tensor(2, dtype=int32, loc=cpu:0, shape=())
- state_dict() Dict[str, Tensor] ¶
Returns a dictionary mapping names to parameters in the module. This will recurse over any nested child modules.
- Returns:
A dictionary mapping names to parameters.
- Return type:
Dict[str, Tensor]
Example
1class MyModule(tp.Module): 2 def __init__(self): 3 super().__init__() 4 self.param = tp.ones((2,), dtype=tp.float32) 5 self.linear1 = tp.Linear(2, 2) 6 self.linear2 = tp.Linear(2, 2) 7 8 9module = MyModule() 10 11state_dict = module.state_dict()
Local Variables¶>>> state_dict { param: tensor([1, 1], dtype=float32, loc=gpu:0, shape=(2,)), linear1.weight: <nvtripy.frontend.module.parameter.DefaultParameter object at 0x79774bc82640>, linear1.bias: <nvtripy.frontend.module.parameter.DefaultParameter object at 0x79774b9cbdf0>, linear2.weight: <nvtripy.frontend.module.parameter.DefaultParameter object at 0x79774b9cfb50>, linear2.bias: <nvtripy.frontend.module.parameter.DefaultParameter object at 0x79774b9cf730>, }
- normalized_shape: Sequence[int]¶
Defines the shape of the input tensor that is to be normalized over.
- eps: float¶
A value added to the denominator to prevent division by zero.