Sequential¶
- class nvtripy.Sequential(*modules: Module | Callable[[Tensor], Tensor] | Dict[str, Module | Callable[[Tensor], Tensor]])[source]¶
Bases:
Module
A module to stack multiple callable layers or modules in a sequential order. The Sequential container can accept either a list of modules/callable objects or a dictionary of named modules/callable objects. Layers are added in the order they are passed, and each is called sequentially during the forward pass.
- Parameters:
*modules (Module | Callable[[Tensor], Tensor] | Dict[str, Module | Callable[[Tensor], Tensor]]) – The module(s) or callable(s) to include in the sequence. These must take exactly one input and return exactly one output. Can be passed as individual positional arguments or as a single dictionary of named modules.
Example: Sequential with Positional Arguments
1model = tp.Sequential(tp.Linear(1, 3), tp.Linear(3, 2)) 2 3model.load_state_dict( 4 { 5 "0.weight": tp.ones((3, 1)), 6 "0.bias": tp.ones((3,)), 7 "1.weight": tp.ones((2, 3)), 8 "1.bias": tp.ones((2,)), 9 } 10) 11 12input = tp.Tensor([1.0]) 13output = model(input)
Local Variables¶>>> model Sequential( 0: Module = Linear( weight: Parameter = (shape=(3, 1), dtype=float32), bias: Parameter = (shape=(3,), dtype=float32), ), 1: Module = Linear( weight: Parameter = (shape=(2, 3), dtype=float32), bias: Parameter = (shape=(2,), dtype=float32), ), ) >>> model.state_dict() { 0.weight: tensor( [[1], [1], [1]], dtype=float32, loc=gpu:0, shape=(3, 1)), 0.bias: tensor([1, 1, 1], dtype=float32, loc=gpu:0, shape=(3,)), 1.weight: tensor( [[1, 1, 1], [1, 1, 1]], dtype=float32, loc=gpu:0, shape=(2, 3)), 1.bias: tensor([1, 1], dtype=float32, loc=gpu:0, shape=(2,)), } >>> input tensor([1], dtype=float32, loc=cpu:0, shape=(1,)) >>> output tensor( [[7, 7]], dtype=float32, loc=gpu:0, shape=(1, 2))
Example: Sequential with a Dictionary
1model = tp.Sequential({"layer1": tp.Linear(1, 3), "layer2": tp.Linear(3, 2)}) 2 3model.load_state_dict( 4 { 5 "layer1.weight": tp.ones((3, 1)), 6 "layer1.bias": tp.ones((3,)), 7 "layer2.weight": tp.ones((2, 3)), 8 "layer2.bias": tp.ones((2,)), 9 } 10) 11 12input = tp.Tensor([1.0]) 13output = model(input)
Local Variables¶>>> model Sequential( layer1: Module = Linear( weight: Parameter = (shape=(3, 1), dtype=float32), bias: Parameter = (shape=(3,), dtype=float32), ), layer2: Module = Linear( weight: Parameter = (shape=(2, 3), dtype=float32), bias: Parameter = (shape=(2,), dtype=float32), ), ) >>> model.state_dict() { layer1.weight: tensor( [[1], [1], [1]], dtype=float32, loc=gpu:0, shape=(3, 1)), layer1.bias: tensor([1, 1, 1], dtype=float32, loc=gpu:0, shape=(3,)), layer2.weight: tensor( [[1, 1, 1], [1, 1, 1]], dtype=float32, loc=gpu:0, shape=(2, 3)), layer2.bias: tensor([1, 1], dtype=float32, loc=gpu:0, shape=(2,)), } >>> input tensor([1], dtype=float32, loc=cpu:0, shape=(1,)) >>> output tensor( [[7, 7]], dtype=float32, loc=gpu:0, shape=(1, 2))
Example: Sequential with Callables
1model = tp.Sequential( 2 tp.relu, lambda x: tp.avgpool(x, kernel_dims=(2, 2), stride=(1, 1)) 3) 4 5input = tp.ones((1, 2, 2, 2), dtype=tp.float32) 6output = model(input)
Local Variables¶>>> model Sequential( ) >>> model.state_dict() {} >>> input tensor( [[[[1, 1], [1, 1]], [[1, 1], [1, 1]]]], dtype=float32, loc=gpu:0, shape=(1, 2, 2, 2)) >>> output tensor( [[[[1]], [[1]]]], dtype=float32, loc=gpu:0, shape=(1, 2, 1, 1))
- forward(input: Tensor) Tensor [source]¶
Defines the forward pass by applying each module in the container sequentially to input
- __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]] [source]¶
Returns an iterator over all the first-order modules in this Sequential container. Each child module is represented by its name and the module object itself.
- Returns:
An iterator over tuples containing the name and module of each child.
- Return type:
Iterator[Tuple[str, Module]]
Example
1model = tp.Sequential(tp.Linear(1, 3), tp.Linear(3, 2)) 2 3for name, child in model.named_children(): 4 print(f"{name}: {type(child).__name__}")
Local Variables¶>>> model Sequential( 0: Module = Linear( weight: Parameter = (shape=(3, 1), dtype=float32), bias: Parameter = (shape=(3,), dtype=float32), ), 1: Module = Linear( weight: Parameter = (shape=(2, 3), dtype=float32), bias: Parameter = (shape=(2,), dtype=float32), ), ) >>> model.state_dict() { 0.weight: <nvtripy.frontend.module.parameter.DefaultParameter object at 0x79774c191c40>, 0.bias: <nvtripy.frontend.module.parameter.DefaultParameter object at 0x79774c191d90>, 1.weight: <nvtripy.frontend.module.parameter.DefaultParameter object at 0x79774c0d6e20>, 1.bias: <nvtripy.frontend.module.parameter.DefaultParameter object at 0x79774c0d6970>, } >>> child Linear( weight: Parameter = (shape=(2, 3), dtype=float32), bias: Parameter = (shape=(2,), dtype=float32), ) >>> child.state_dict() { weight: <nvtripy.frontend.module.parameter.DefaultParameter object at 0x79774c0d6e20>, bias: <nvtripy.frontend.module.parameter.DefaultParameter object at 0x79774c0d6970>, }
Output¶0: Linear 1: 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 0x79774c0d61f0>, linear1.bias: <nvtripy.frontend.module.parameter.DefaultParameter object at 0x79774c183c70>, linear2.weight: <nvtripy.frontend.module.parameter.DefaultParameter object at 0x79774c1fb250>, linear2.bias: <nvtripy.frontend.module.parameter.DefaultParameter object at 0x79774c1fb370>, }