Module

class nvtripy.Module[source]

Bases: object

Base class used to define neural network modules. Child classes must implement the forward() method.

You can nest modules by assigning them as attributes of other modules.

Child modules, nvtripy.Tensor s, or other callables/lambda functions may be contained in Python lists or dicts.

If using dicts, the keys must be strings. Nested data structures (for example, lists of lists) are not supported. Taking child modules as an example, this is allowed:

self.linear = tp.Linear(2, 2)
self.list_modules = [tp.Linear(2, 2), tp.Linear(2, 2)]
self.dict_modules = {
    "linear": tp.Linear(2, 2),
    "layernorm": tp.LayerNorm(2),
}

This is another valid example with a wrapped nvtripy.avgpool lambda function

self.dict_modules = {
    "convolution": tp.Conv(in_channels=2, out_channels=2, kernel_dims=(1,1)),
    "pool": lambda x: tp.avgpool(x, kernel_dims=(2,2))
}

Whereas this is not supported:

self.list_modules = [[tp.Linear(2, 2)], [tp.Linear(2, 2)]]
self.dict_modules = {
    (1, "linear"): tp.Linear(2, 2),
}
Example
 1class AddBias(tp.Module):
 2    def __init__(self):
 3        super().__init__()
 4        self.bias = tp.Tensor([1.0, 1.0], dtype=tp.float32)
 5
 6    def forward(self, x):
 7        return x + self.bias
 8
 9
10add_bias = AddBias()
11
12input = tp.Tensor([1.0, 1.0], dtype=tp.float32)
13output = add_bias(input)
Local Variables
>>> add_bias
AddBias(
    bias: Parameter = (shape=(2,), dtype=float32),
)
>>> add_bias.state_dict()
{
    bias: tensor([1, 1], dtype=float32, loc=cpu:0, shape=(2,)),
}

>>> input
tensor([1, 1], dtype=float32, loc=cpu:0, shape=(2,))

>>> output
tensor([2, 2], dtype=float32, loc=gpu:0, shape=(2,))
state_dict() → Dict[str, Tensor][source]

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 0x79774c315bb0>,
    linear1.bias: <nvtripy.frontend.module.parameter.DefaultParameter object at 0x79774c327be0>,
    linear2.weight: <nvtripy.frontend.module.parameter.DefaultParameter object at 0x79774c330580>,
    linear2.bias: <nvtripy.frontend.module.parameter.DefaultParameter object at 0x79774c330670>,
}
load_state_dict(state_dict: Dict[str, Tensor], strict: bool = True) → Tuple[Set[str], Set[str]][source]

Loads parameters from the provided state_dict into the current module. This will recurse over any nested child modules.

Parameters:

  • state_dict (Dict[str, Tensor]) – A dictionary mapping names to parameters.

  • strict (bool) – If True, keys in state_dict must exactly match those in this module. If not, an error will be raised.

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 two sets 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

state_dict()

named_children() → Iterator[Tuple[str, Module]][source]

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]][source]
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=())
__call__(*args: Any, **kwargs: Any) → Any[source]

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,))
abstract forward(*args: Any, **kwargs: Any) → Any[source]

Implements the forward pass of this module. Child classes must implement this method.

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,))

See also: