Module

class tripy.Module[source]

Bases: object

Base class used to define neural network modules. You can nest modules by assigning them as attributes of other modules.

Child modules or tripy.Parameter s 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),
}

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
Example
 1class AddBias(tp.Module):
 2    def __init__(self):
 3        super().__init__()
 4        self.bias = tp.Parameter(tp.Tensor([1.0, 1.0], dtype=tp.float32))
 5
 6    def __call__(self, x):
 7        return x + self.bias
 8
 9add_bias = AddBias()
10
11input = tp.Tensor([1.0, 1.0], dtype=tp.float32)
12output = add_bias(input)

>>> 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, Parameter][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, Parameter]

Example
Example
 1class MyModule(tp.Module):
 2    def __init__(self):
 3        super().__init__()
 4        self.param = tp.Parameter(tp.ones((2,), dtype=tp.float32))
 5        self.linear1 = tp.Linear(2, 2)
 6        self.linear2 = tp.Linear(2, 2)
 7
 8module = MyModule()
 9
10state_dict = module.state_dict()

>>> 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, Parameter], 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, Parameter]) – 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
Example
1# Using the `module` and `state_dict` from the `state_dict()` example:
2print(f"Before: {module.param}")
3
4state_dict["param"] = tp.Parameter(tp.zeros((2,), dtype=tp.float32))
5module.load_state_dict(state_dict)
6
7print(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]][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
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
 7stacked_linear = StackedLinear()
 8
 9for name, module in stacked_linear.named_children():
10    print(f"{name}: {type(module).__name__}")
Output:
linear1: Linear
linear2: Linear
named_parameters() → Iterator[Tuple[str, Parameter]][source]
Returns:

An iterator over tuples containing the name of a parameter and the parameter itself.

Return type:

Iterator[Tuple[str, Parameter]]

Example
Example
 1class Linear(tp.Module):
 2    def __init__(self):
 3        super().__init__()
 4        self.alpha = tp.Parameter(1)
 5        self.beta = tp.Parameter(2)
 6
 7linear = Linear()
 8
 9for name, parameter in linear.named_parameters():
10    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: