reference/generated/tensordict.from_modules

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

from_modules

tensordict.from_modules(*modules, as_module: bool = False, lock: bool = True, use_state_dict: bool = False, lazy_stack: bool = False, expand_identical: bool = False)

Stacks the parameters and buffers of several modules in one tensordict.

The result has one batch dimension, with one element per module. Use it to run an ensemble of modules that have the same structure with torch.vmap() and to_module(). Same as tensordict.TensorDict.from_modules(): see that page for the difference between a dense and a lazy result, and for a full ensemble example.

Parameters:

*modules (Module) – the modules to get the parameters from.

Keyword Arguments:

• as_module (bool, optional) – if True, a TensorDictParams instance is returned. Defaults to False.

• lock (bool, optional) – if True, the resulting tensordict is locked. Defaults to True.

• use_state_dict (bool, optional) – if True, the state-dict of each module is used. Defaults to False.

• lazy_stack (bool, optional) – if True, the parameters are stacked in a LazyStackedTensorDict that keeps the original parameters. If False, they are stacked into new parameters. Exclusive with as_module. Defaults to False.

• expand_identical (bool, optional) – if True, a parameter that is the same object in every module is expanded instead of stacked. Ignored when lazy_stack=True. Defaults to False.

Returns:

A TensorDict with batch size [len(modules)], a LazyStackedTensorDict if lazy_stack=True, or a TensorDictParams if as_module=True.

Examples

>>> import torch
>>> from torch import nn
>>> import tensordict
>>> modules = [nn.Linear(3, 4) for _ in range(2)]
>>> params = tensordict.from_modules(*modules)
>>> params.batch_size
torch.Size([2])
>>> params["weight"].shape
torch.Size([2, 4, 3])