Module

class nvtripy.Module

Bases: object

Base class used to define neural network modules. 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

class AddBias(tp.Module):
    def __init__(self):
        super().__init__()
        self.bias = tp.Tensor([1.0, 1.0], dtype=tp.float32)

    def __call__(self, x):
        return x + self.bias


add_bias = AddBias()

input = tp.Tensor([1.0, 1.0], dtype=tp.float32)
output = add_bias(input)

Local Variables

>>> add_bias
AddBias(
    bias: Parameter = (shape=[2], dtype=float32),
)
>>> add_bias.state_dict()
{
    bias: tensor([1.0000, 1.0000], dtype=float32, loc=gpu:0, shape=(2,)),
}

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

>>> output
tensor([2.0000, 2.0000], dtype=float32, loc=gpu:0, shape=(2,))

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

class MyModule(tp.Module):
    def __init__(self):
        super().__init__()
        self.param = tp.ones((2,), dtype=tp.float32)
        self.linear1 = tp.Linear(2, 2)
        self.linear2 = tp.Linear(2, 2)


module = MyModule()

state_dict = module.state_dict()

Local Variables

>>> state_dict
{
    param: tensor([1.0000, 1.0000], dtype=float32, loc=gpu:0, shape=(2,)),
    linear1.weight: tensor(
        [[0.0000, 1.0000],
         [2.0000, 3.0000]], 
        dtype=float32, loc=gpu:0, shape=(2, 2)),
    linear1.bias: tensor([0.0000, 1.0000], dtype=float32, loc=gpu:0, shape=(2,)),
    linear2.weight: tensor(
        [[0.0000, 1.0000],
         [2.0000, 3.0000]], 
        dtype=float32, loc=gpu:0, shape=(2, 2)),
    linear2.bias: tensor([0.0000, 1.0000], dtype=float32, loc=gpu:0, shape=(2,)),
}

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:

• 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

# Using the `module` and `state_dict` from the `state_dict()` example:
print(f"Before: {module.param}")

state_dict["param"] = tp.zeros((2,), dtype=tp.float32)
module.load_state_dict(state_dict)

print(f"After: {module.param}")

Output

Before: tensor([1.0000, 1.0000], dtype=float32, loc=gpu:0, shape=(2,))
After: tensor([0.0000, 0.0000], dtype=float32, loc=gpu:0, shape=(2,))

See also

state_dict()

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

class StackedLinear(tp.Module):
    def __init__(self):
        super().__init__()
        self.linear1 = tp.Linear(2, 2)
        self.linear2 = tp.Linear(2, 2)


stacked_linear = StackedLinear()

for name, module in stacked_linear.named_children():
    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

class MyModule(tp.Module):
    def __init__(self):
        super().__init__()
        self.alpha = tp.Tensor(1)
        self.beta = tp.Tensor(2)


linear = MyModule()

for name, parameter in linear.named_parameters():
    print(f"{name}: {parameter}")

Output

alpha: tensor(1, dtype=int32, loc=gpu:0, shape=())
beta: tensor(2, dtype=int32, loc=gpu:0, shape=())

See also:

• Sequential