Rate this Page

★★★★★

Module #

classtorch.nn.Module(*args,**kwargs)[source]#

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested ina tree structure. You can assign the submodules as regular attributes:

importtorch.nnasnnimporttorch.nn.functionalasFclassModel(nn.Module):def__init__(self)->None:super().__init__()self.conv1=nn.Conv2d(1,20,5)self.conv2=nn.Conv2d(20,20,5)defforward(self,x):x=F.relu(self.conv1(x))returnF.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have theirparameters converted when you callto(), etc.

Note

As per the example above, an__init__() call to the parent classmust be made before assignment on the child.

Variables: training (bool) – Boolean represents whether this module is in training orevaluation mode.

add_module(name,module)[source]#

Add a child module to the current module.

The module can be accessed as an attribute using the given name.

Parameters

name (str) – name of the child module. The child module can beaccessed from this module using the given name
module (Module) – child module to be added to the module.

apply(fn)[source]#

Applyfn recursively to every submodule (as returned by.children()) as well as self.

Typical use includes initializing the parameters of a model(see alsotorch.nn.init).

Parameters: fn (Module -> None) – function to be applied to each submodule
Returns: self
Return type: Module

Example:

>>>@torch.no_grad()>>>definit_weights(m):>>>print(m)>>>iftype(m)==nn.Linear:>>>m.weight.fill_(1.0)>>>print(m.weight)>>>net=nn.Sequential(nn.Linear(2,2),nn.Linear(2,2))>>>net.apply(init_weights)Linear(in_features=2, out_features=2, bias=True)Parameter containing:tensor([[1., 1.],        [1., 1.]], requires_grad=True)Linear(in_features=2, out_features=2, bias=True)Parameter containing:tensor([[1., 1.],        [1., 1.]], requires_grad=True)Sequential(  (0): Linear(in_features=2, out_features=2, bias=True)  (1): Linear(in_features=2, out_features=2, bias=True))

bfloat16()[source]#

Casts all floating point parameters and buffers tobfloat16 datatype.

Note

This method modifies the module in-place.

Returns: self
Return type: Module

buffers(recurse=True)[source]#

Return an iterator over module buffers.

Parameters: recurse (bool) – if True, then yields buffers of this moduleand all submodules. Otherwise, yields only buffers thatare direct members of this module.
Yields: torch.Tensor – module buffer
Return type: Iterator[Tensor]

Example:

>>>forbufinmodel.buffers():>>>print(type(buf),buf.size())<class 'torch.Tensor'> (20L,)<class 'torch.Tensor'> (20L, 1L, 5L, 5L)

children()[source]#

Return an iterator over immediate children modules.

Yields: Module – a child module
Return type: Iterator[Module]

compile(*args,**kwargs)[source]#

Compile this Module’s forward usingtorch.compile().

This Module’s__call__ method is compiled and all arguments are passed as-istotorch.compile().

Seetorch.compile() for details on the arguments for this function.

cpu()[source]#

Move all model parameters and buffers to the CPU.

Note

This method modifies the module in-place.

Returns: self
Return type: Module

cuda(device=None)[source]#

Move all model parameters and buffers to the GPU.

This also makes associated parameters and buffers different objects. Soit should be called before constructing the optimizer if the module willlive on GPU while being optimized.

Note

This method modifies the module in-place.

Parameters: device (int,optional) – if specified, all parameters will becopied to that device
Returns: self
Return type: Module

double()[source]#

Casts all floating point parameters and buffers todouble datatype.

Note

This method modifies the module in-place.

Returns: self
Return type: Module

eval()[source]#

Set the module in evaluation mode.

This has an effect only on certain modules. See the documentation ofparticular modules for details of their behaviors in training/evaluationmode, i.e. whether they are affected, e.g.Dropout,BatchNorm,etc.

This is equivalent withself.train(False).

SeeLocally disabling gradient computation for a comparison between.eval() and several similar mechanisms that may be confused with it.

Returns: self
Return type: Module

extra_repr()[source]#

Return the extra representation of the module.

To print customized extra information, you should re-implementthis method in your own modules. Both single-line and multi-linestrings are acceptable.

Return type: str

float()[source]#

Casts all floating point parameters and buffers tofloat datatype.

Note

This method modifies the module in-place.

Returns: self
Return type: Module

forward(*input)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined withinthis function, one should call theModule instance afterwardsinstead of this since the former takes care of running theregistered hooks while the latter silently ignores them.

get_buffer(target)[source]#

Return the buffer given bytarget if it exists, otherwise throw an error.

See the docstring forget_submodule for a more detailedexplanation of this method’s functionality as well as how tocorrectly specifytarget.

Parameters: target (str) – The fully-qualified string name of the bufferto look for. (Seeget_submodule for how to specify afully-qualified string.)
Returns: The buffer referenced bytarget
Return type: torch.Tensor
Raises: AttributeError – If the target string references an invalid path or resolves to something that is not a buffer

get_extra_state()[source]#

Return any extra state to include in the module’s state_dict.

Implement this and a correspondingset_extra_state() for your moduleif you need to store extra state. This function is called when building themodule’sstate_dict().

Note that extra state should be picklable to ensure working serializationof the state_dict. We only provide backwards compatibility guaranteesfor serializing Tensors; other objects may break backwards compatibility iftheir serialized pickled form changes.

Returns: Any extra state to store in the module’s state_dict
Return type: object

get_parameter(target)[source]#

Return the parameter given bytarget if it exists, otherwise throw an error.

See the docstring forget_submodule for a more detailedexplanation of this method’s functionality as well as how tocorrectly specifytarget.

Parameters: target (str) – The fully-qualified string name of the Parameterto look for. (Seeget_submodule for how to specify afully-qualified string.)
Returns: The Parameter referenced bytarget
Return type: torch.nn.Parameter
Raises: AttributeError – If the target string references an invalid path or resolves to something that is not annn.Parameter

get_submodule(target)[source]#

Return the submodule given bytarget if it exists, otherwise throw an error.

For example, let’s say you have annn.ModuleA thatlooks like this:

A(    (net_b): Module(        (net_c): Module(            (conv): Conv2d(16, 33, kernel_size=(3, 3), stride=(2, 2))        )        (linear): Linear(in_features=100, out_features=200, bias=True)    ))

(The diagram shows annn.ModuleA.A which has a nestedsubmodulenet_b, which itself has two submodulesnet_candlinear.net_c then has a submoduleconv.)

To check whether or not we have thelinear submodule, wewould callget_submodule("net_b.linear"). To check whetherwe have theconv submodule, we would callget_submodule("net_b.net_c.conv").

The runtime ofget_submodule is bounded by the degreeof module nesting intarget. A query againstnamed_modules achieves the same result, but it is O(N) inthe number of transitive modules. So, for a simple check to seeif some submodule exists,get_submodule should always beused.

Parameters: target (str) – The fully-qualified string name of the submoduleto look for. (See above example for how to specify afully-qualified string.)
Returns: The submodule referenced bytarget
Return type: torch.nn.Module
Raises: AttributeError – If at any point along the path resulting from the target string the (sub)path resolves to a non-existent attribute name or an object that is not an instance ofnn.Module.

half()[source]#

Casts all floating point parameters and buffers tohalf datatype.

Note

This method modifies the module in-place.

Returns: self
Return type: Module

ipu(device=None)[source]#

Move all model parameters and buffers to the IPU.

This also makes associated parameters and buffers different objects. Soit should be called before constructing the optimizer if the module willlive on IPU while being optimized.

Note

This method modifies the module in-place.

Parameters: device (int,optional) – if specified, all parameters will becopied to that device
Returns: self
Return type: Module

load_state_dict(state_dict,strict=True,assign=False)[source]#

Copy parameters and buffers fromstate_dict into this module and its descendants.

Ifstrict isTrue, thenthe keys ofstate_dict must exactly match the keys returnedby this module’sstate_dict() function.

Warning

Ifassign isTrue the optimizer must be created afterthe call toload_state_dict unlessget_swap_module_params_on_conversion() isTrue.

Parameters

state_dict (dict) – a dict containing parameters andpersistent buffers.
strict (bool,optional) – whether to strictly enforce that the keysinstate_dict match the keys returned by this module’sstate_dict() function. Default:True
assign (bool,optional) – When set toFalse, the properties of the tensorsin the current module are preserved whereas setting it toTrue preservesproperties of the Tensors in the state dict. The onlyexception is therequires_grad field ofParameterfor which the value from the module is preserved. Default:False

Returns

missing_keys is a list of str containing any keys that are expected
by this module but missing from the providedstate_dict.
unexpected_keys is a list of str containing the keys that are not
expected by this module but present in the providedstate_dict.

Return type

NamedTuple withmissing_keys andunexpected_keys fields

Note

If a parameter or buffer is registered asNone and its corresponding keyexists instate_dict,load_state_dict() will raise aRuntimeError.

modules()[source]#

Return an iterator over all modules in the network.

Yields: Module – a module in the network
Return type: Iterator[Module]

Note

Duplicate modules are returned only once. In the followingexample,l will be returned only once.

Example:

>>>l=nn.Linear(2,2)>>>net=nn.Sequential(l,l)>>>foridx,minenumerate(net.modules()):...print(idx,'->',m)0 -> Sequential(  (0): Linear(in_features=2, out_features=2, bias=True)  (1): Linear(in_features=2, out_features=2, bias=True))1 -> Linear(in_features=2, out_features=2, bias=True)

mtia(device=None)[source]#

Move all model parameters and buffers to the MTIA.

This also makes associated parameters and buffers different objects. Soit should be called before constructing the optimizer if the module willlive on MTIA while being optimized.

Note

This method modifies the module in-place.

Parameters: device (int,optional) – if specified, all parameters will becopied to that device
Returns: self
Return type: Module

named_buffers(prefix='',recurse=True,remove_duplicate=True)[source]#

Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.

Parameters

prefix (str) – prefix to prepend to all buffer names.
recurse (bool,optional) – if True, then yields buffers of this moduleand all submodules. Otherwise, yields only buffers thatare direct members of this module. Defaults to True.
remove_duplicate (bool,optional) – whether to remove the duplicated buffers in the result. Defaults to True.

Yields

(str, torch.Tensor) – Tuple containing the name and buffer

Return type

Iterator[tuple[str,torch.Tensor]]

Example:

>>>forname,bufinself.named_buffers():>>>ifnamein['running_var']:>>>print(buf.size())

named_children()[source]#

Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.

Yields: (str, Module) – Tuple containing a name and child module
Return type: Iterator[tuple[str, ‘Module’]]

Example:

>>>forname,moduleinmodel.named_children():>>>ifnamein['conv4','conv5']:>>>print(module)

named_modules(memo=None,prefix='',remove_duplicate=True)[source]#

Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.

Parameters

memo (Optional[set['Module']]) – a memo to store the set of modules already added to the result
prefix (str) – a prefix that will be added to the name of the module
remove_duplicate (bool) – whether to remove the duplicated module instances in the resultor not

Yields

(str, Module) – Tuple of name and module

Note

Duplicate modules are returned only once. In the followingexample,l will be returned only once.

Example:

>>>l=nn.Linear(2,2)>>>net=nn.Sequential(l,l)>>>foridx,minenumerate(net.named_modules()):...print(idx,'->',m)0 -> ('', Sequential(  (0): Linear(in_features=2, out_features=2, bias=True)  (1): Linear(in_features=2, out_features=2, bias=True)))1 -> ('0', Linear(in_features=2, out_features=2, bias=True))

named_parameters(prefix='',recurse=True,remove_duplicate=True)[source]#

Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.

Parameters

prefix (str) – prefix to prepend to all parameter names.
recurse (bool) – if True, then yields parameters of this moduleand all submodules. Otherwise, yields only parameters thatare direct members of this module.
remove_duplicate (bool,optional) – whether to remove the duplicatedparameters in the result. Defaults to True.

Yields

(str, Parameter) – Tuple containing the name and parameter

Return type

Iterator[tuple[str,torch.nn.parameter.Parameter]]

Example:

>>>forname,paraminself.named_parameters():>>>ifnamein['bias']:>>>print(param.size())

parameters(recurse=True)[source]#

Return an iterator over module parameters.

This is typically passed to an optimizer.

Parameters: recurse (bool) – if True, then yields parameters of this moduleand all submodules. Otherwise, yields only parameters thatare direct members of this module.
Yields: Parameter – module parameter
Return type: Iterator[Parameter]

Example:

>>>forparaminmodel.parameters():>>>print(type(param),param.size())<class 'torch.Tensor'> (20L,)<class 'torch.Tensor'> (20L, 1L, 5L, 5L)

register_backward_hook(hook)[source]#

This function is deprecated in favor ofregister_full_backward_hook() andthe behavior of this function will change in future versions.

Returns: a handle that can be used to remove the added hook by callinghandle.remove()
Return type: torch.utils.hooks.RemovableHandle

register_buffer(name,tensor,persistent=True)[source]#

Add a buffer to the module.

This is typically used to register a buffer that should not beconsidered a model parameter. For example, BatchNorm’srunning_meanis not a parameter, but is part of the module’s state. Buffers, bydefault, are persistent and will be saved alongside parameters. Thisbehavior can be changed by settingpersistent toFalse. Theonly difference between a persistent buffer and a non-persistent bufferis that the latter will not be a part of this module’sstate_dict.

Buffers can be accessed as attributes using given names.

Parameters

name (str) – name of the buffer. The buffer can be accessedfrom this module using the given name
tensor (Tensor orNone) – buffer to be registered. IfNone, then operationsthat run on buffers, such ascuda, are ignored. IfNone,the buffer isnot included in the module’sstate_dict.
persistent (bool) – whether the buffer is part of this module’sstate_dict.

Example:

>>>self.register_buffer('running_mean',torch.zeros(num_features))

register_forward_hook(hook,*,prepend=False,with_kwargs=False,always_call=False)[source]#

The hook will be called every time afterforward() has computed an output.

Ifwith_kwargs isFalse or not specified, the input contains onlythe positional arguments given to the module. Keyword arguments won’t bepassed to the hooks and only to theforward. The hook can modify theoutput. It can modify the input inplace but it will not have effect onforward since this is called afterforward() is called. The hookshould have the following signature:

hook(module,args,output)->Noneormodifiedoutput

Ifwith_kwargs isTrue, the forward hook will be passed thekwargs given to the forward function and be expected to return theoutput possibly modified. The hook should have the following signature:

hook(module,args,kwargs,output)->Noneormodifiedoutput

Parameters

hook (Callable) – The user defined hook to be registered.
prepend (bool) – IfTrue, the providedhook will be firedbefore all existingforward hooks on thistorch.nn.Module. Otherwise, the providedhook will be fired after all existingforward hooks onthistorch.nn.Module. Note that globalforward hooks registered withregister_module_forward_hook() will fire before all hooksregistered by this method.Default:False
with_kwargs (bool) – IfTrue, thehook will be passed thekwargs given to the forward function.Default:False
always_call (bool) – IfTrue thehook will be run regardless ofwhether an exception is raised while calling the Module.Default:False

Returns

a handle that can be used to remove the added hook by callinghandle.remove()

Return type

torch.utils.hooks.RemovableHandle

register_forward_pre_hook(hook,*,prepend=False,with_kwargs=False)[source]#

The hook will be called every time beforeforward() is invoked.

Ifwith_kwargs is false or not specified, the input contains onlythe positional arguments given to the module. Keyword arguments won’t bepassed to the hooks and only to theforward. The hook can modify theinput. User can either return a tuple or a single modified value in thehook. We will wrap the value into a tuple if a single value is returned(unless that value is already a tuple). The hook should have thefollowing signature:

hook(module,args)->Noneormodifiedinput

Ifwith_kwargs is true, the forward pre-hook will be passed thekwargs given to the forward function. And if the hook modifies theinput, both the args and kwargs should be returned. The hook should havethe following signature:

hook(module,args,kwargs)->Noneoratupleofmodifiedinputandkwargs

Parameters

hook (Callable) – The user defined hook to be registered.
prepend (bool) – If true, the providedhook will be fired beforeall existingforward_pre hooks on thistorch.nn.Module. Otherwise, the providedhook will be fired after all existingforward_pre hookson thistorch.nn.Module. Note that globalforward_pre hooks registered withregister_module_forward_pre_hook() will fire before allhooks registered by this method.Default:False
with_kwargs (bool) – If true, thehook will be passed the kwargsgiven to the forward function.Default:False

Returns

a handle that can be used to remove the added hook by callinghandle.remove()

Return type

torch.utils.hooks.RemovableHandle

register_full_backward_hook(hook,prepend=False)[source]#

The hook will be called every time the gradients with respect to a module are computed, and its firing rules are as follows:

Ordinarily, the hook fires when the gradients are computed with respect to the module inputs.
If none of the module inputs require gradients, the hook will fire when the gradients are computedwith respect to module outputs.
If none of the module outputs require gradients, then the hooks will not fire.

The hook should have the following signature:

hook(module,grad_input,grad_output)->tuple(Tensor)orNone

Thegrad_input andgrad_output are tuples that contain the gradientswith respect to the inputs and outputs respectively. The hook shouldnot modify its arguments, but it can optionally return a new gradient withrespect to the input that will be used in place ofgrad_input insubsequent computations.grad_input will only correspond to the inputs givenas positional arguments and all kwarg arguments are ignored. Entriesingrad_input andgrad_output will beNone for all non-Tensorarguments.

For technical reasons, when this hook is applied to a Module, its forward function willreceive a view of each Tensor passed to the Module. Similarly the caller will receive a viewof each Tensor returned by the Module’s forward function.

Warning

Modifying inputs or outputs inplace is not allowed when using backward hooks andwill raise an error.

Parameters

hook (Callable) – The user-defined hook to be registered.
prepend (bool) – If true, the providedhook will be fired beforeall existingbackward hooks on thistorch.nn.Module. Otherwise, the providedhook will be fired after all existingbackward hooks onthistorch.nn.Module. Note that globalbackward hooks registered withregister_module_full_backward_hook() will fire beforeall hooks registered by this method.

Returns

a handle that can be used to remove the added hook by callinghandle.remove()

Return type

torch.utils.hooks.RemovableHandle

register_full_backward_pre_hook(hook,prepend=False)[source]#

The hook will be called every time the gradients for the module are computed.The hook should have the following signature:

hook(module,grad_output)->tuple[Tensor]orNone

Thegrad_output is a tuple. The hook shouldnot modify its arguments, but it can optionally return a new gradient withrespect to the output that will be used in place ofgrad_output insubsequent computations. Entries ingrad_output will beNone forall non-Tensor arguments.

Warning

Modifying inputs inplace is not allowed when using backward hooks andwill raise an error.

Parameters

hook (Callable) – The user-defined hook to be registered.
prepend (bool) – If true, the providedhook will be fired beforeall existingbackward_pre hooks on thistorch.nn.Module. Otherwise, the providedhook will be fired after all existingbackward_pre hookson thistorch.nn.Module. Note that globalbackward_pre hooks registered withregister_module_full_backward_pre_hook() will fire beforeall hooks registered by this method.

Returns

a handle that can be used to remove the added hook by callinghandle.remove()

Return type

torch.utils.hooks.RemovableHandle

register_load_state_dict_post_hook(hook)[source]#

It should have the following signature::: hook(module, incompatible_keys) -> None

Themodule argument is the current module that this hook is registeredon, and theincompatible_keys argument is aNamedTuple consistingof attributesmissing_keys andunexpected_keys.missing_keysis alist ofstr containing the missing keys andunexpected_keys is alist ofstr containing the unexpected keys.

The given incompatible_keys can be modified inplace if needed.

Note that the checks performed when callingload_state_dict() withstrict=True are affected by modifications the hook makes tomissing_keys orunexpected_keys, as expected. Additions to eitherset of keys will result in an error being thrown whenstrict=True, andclearing out both missing and unexpected keys will avoid an error.

Returns: a handle that can be used to remove the added hook by callinghandle.remove()
Return type: torch.utils.hooks.RemovableHandle

register_load_state_dict_pre_hook(hook)[source]#

It should have the following signature::: hook(module, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs) -> None # noqa: B950

Parameters: hook (Callable) – Callable hook that will be invoked beforeloading the state dict.

register_module(name,module)[source]#

Alias foradd_module().

register_parameter(name,param)[source]#

Add a parameter to the module.

The parameter can be accessed as an attribute using given name.

Parameters

name (str) – name of the parameter. The parameter can be accessedfrom this module using the given name
param (Parameter orNone) – parameter to be added to the module. IfNone, then operations that run on parameters, such ascuda,are ignored. IfNone, the parameter isnot included in themodule’sstate_dict.

register_state_dict_post_hook(hook)[source]#

It should have the following signature::: hook(module, state_dict, prefix, local_metadata) -> None

The registered hooks can modify thestate_dict inplace.

register_state_dict_pre_hook(hook)[source]#

It should have the following signature::: hook(module, prefix, keep_vars) -> None

The registered hooks can be used to perform pre-processing before thestate_dictcall is made.

requires_grad_(requires_grad=True)[source]#

Change if autograd should record operations on parameters in this module.

This method sets the parameters’requires_grad attributesin-place.

This method is helpful for freezing part of the module for finetuningor training parts of a model individually (e.g., GAN training).

SeeLocally disabling gradient computation for a comparison between.requires_grad_() and several similar mechanisms that may be confused with it.

Parameters: requires_grad (bool) – whether autograd should record operations onparameters in this module. Default:True.
Returns: self
Return type: Module

set_extra_state(state)[source]#

Set extra state contained in the loadedstate_dict.

This function is called fromload_state_dict() to handle any extra statefound within thestate_dict. Implement this function and a correspondingget_extra_state() for your module if you need to store extra state within itsstate_dict.

Parameters: state (dict) – Extra state from thestate_dict

set_submodule(target,module,strict=False)[source]#

Set the submodule given bytarget if it exists, otherwise throw an error.

Note

Ifstrict is set toFalse (default), the method will replace an existing submoduleor create a new submodule if the parent module exists. Ifstrict is set toTrue,the method will only attempt to replace an existing submodule and throw an error ifthe submodule does not exist.

For example, let’s say you have annn.ModuleA thatlooks like this:

A(    (net_b): Module(        (net_c): Module(            (conv): Conv2d(3, 3, 3)        )        (linear): Linear(3, 3)    ))

(The diagram shows annn.ModuleA.A has a nestedsubmodulenet_b, which itself has two submodulesnet_candlinear.net_c then has a submoduleconv.)

To override theConv2d with a new submoduleLinear, youcould callset_submodule("net_b.net_c.conv",nn.Linear(1,1))wherestrict could beTrue orFalse

To add a new submoduleConv2d to the existingnet_b module,you would callset_submodule("net_b.conv",nn.Conv2d(1,1,1)).

In the above if you setstrict=True and callset_submodule("net_b.conv",nn.Conv2d(1,1,1),strict=True), an AttributeErrorwill be raised becausenet_b does not have a submodule namedconv.

Parameters

target (str) – The fully-qualified string name of the submoduleto look for. (See above example for how to specify afully-qualified string.)
module (Module) – The module to set the submodule to.
strict (bool) – IfFalse, the method will replace an existing submoduleor create a new submodule if the parent module exists. IfTrue,the method will only attempt to replace an existing submodule and throw an errorif the submodule doesn’t already exist.

Raises

ValueError – If thetarget string is empty or ifmodule is not an instance ofnn.Module.
AttributeError – If at any point along the path resulting from thetarget string the (sub)path resolves to a non-existent attribute name or an object that is not an instance ofnn.Module.

share_memory()[source]#

Seetorch.Tensor.share_memory_().

Return type: Self

state_dict(*,destination:T_destination,prefix:str='',keep_vars:bool=False)→T_destination[source]#

state_dict(*,prefix:str='',keep_vars:bool=False)→dict[str,Any]

Return a dictionary containing references to the whole state of the module.

Both parameters and persistent buffers (e.g. running averages) areincluded. Keys are corresponding parameter and buffer names.Parameters and buffers set toNone are not included.

Note

The returned object is a shallow copy. It contains referencesto the module’s parameters and buffers.

Warning

Currentlystate_dict() also accepts positional arguments fordestination,prefix andkeep_vars in order. However,this is being deprecated and keyword arguments will be enforced infuture releases.

Warning

Please avoid the use of argumentdestination as it is notdesigned for end-users.

Parameters

destination (dict,optional) – If provided, the state of module willbe updated into the dict and the same object is returned.Otherwise, anOrderedDict will be created and returned.Default:None.
prefix (str,optional) – a prefix added to parameter and buffernames to compose the keys in state_dict. Default:''.
keep_vars (bool,optional) – by default theTensor sreturned in the state dict are detached from autograd. If it’sset toTrue, detaching will not be performed.Default:False.

Returns

a dictionary containing a whole state of the module

Return type

dict

Example:

>>>module.state_dict().keys()['bias', 'weight']

to(device:Optional[Union[str,device,int]]=...,dtype:Optional[dtype]=...,non_blocking:bool=...)→Self[source]#

to(dtype:dtype,non_blocking:bool=...)→Self

to(tensor:Tensor,non_blocking:bool=...)→Self

Move and/or cast the parameters and buffers.

This can be called as

to(device=None,dtype=None,non_blocking=False)[source]

to(dtype,non_blocking=False)[source]

to(tensor,non_blocking=False)[source]

to(memory_format=torch.channels_last)[source]

Its signature is similar totorch.Tensor.to(), but only acceptsfloating point or complexdtypes. In addition, this method willonly cast the floating point or complex parameters and buffers todtype(if given). The integral parameters and buffers will be moveddevice, if that is given, but with dtypes unchanged. Whennon_blocking is set, it tries to convert/move asynchronouslywith respect to the host if possible, e.g., moving CPU Tensors withpinned memory to CUDA devices.

See below for examples.

Note

This method modifies the module in-place.

Parameters

device (torch.device) – the desired device of the parametersand buffers in this module
dtype (torch.dtype) – the desired floating point or complex dtype ofthe parameters and buffers in this module
tensor (torch.Tensor) – Tensor whose dtype and device are the desireddtype and device for all parameters and buffers in this module
memory_format (torch.memory_format) – the desired memoryformat for 4D parameters and buffers in this module (keywordonly argument)

Returns

self

Return type

Module

Examples:

>>>linear=nn.Linear(2,2)>>>linear.weightParameter containing:tensor([[ 0.1913, -0.3420],        [-0.5113, -0.2325]])>>>linear.to(torch.double)Linear(in_features=2, out_features=2, bias=True)>>>linear.weightParameter containing:tensor([[ 0.1913, -0.3420],        [-0.5113, -0.2325]], dtype=torch.float64)>>>gpu1=torch.device("cuda:1")>>>linear.to(gpu1,dtype=torch.half,non_blocking=True)Linear(in_features=2, out_features=2, bias=True)>>>linear.weightParameter containing:tensor([[ 0.1914, -0.3420],        [-0.5112, -0.2324]], dtype=torch.float16, device='cuda:1')>>>cpu=torch.device("cpu")>>>linear.to(cpu)Linear(in_features=2, out_features=2, bias=True)>>>linear.weightParameter containing:tensor([[ 0.1914, -0.3420],        [-0.5112, -0.2324]], dtype=torch.float16)>>>linear=nn.Linear(2,2,bias=None).to(torch.cdouble)>>>linear.weightParameter containing:tensor([[ 0.3741+0.j,  0.2382+0.j],        [ 0.5593+0.j, -0.4443+0.j]], dtype=torch.complex128)>>>linear(torch.ones(3,2,dtype=torch.cdouble))tensor([[0.6122+0.j, 0.1150+0.j],        [0.6122+0.j, 0.1150+0.j],        [0.6122+0.j, 0.1150+0.j]], dtype=torch.complex128)

to_empty(*,device,recurse=True)[source]#

Move the parameters and buffers to the specified device without copying storage.

Parameters

device (torch.device) – The desired device of the parametersand buffers in this module.
recurse (bool) – Whether parameters and buffers of submodules shouldbe recursively moved to the specified device.

Returns

self

Return type

Module

train(mode=True)[source]#

Set the module in training mode.

This has an effect only on certain modules. See the documentation ofparticular modules for details of their behaviors in training/evaluationmode, i.e., whether they are affected, e.g.Dropout,BatchNorm,etc.

Parameters: mode (bool) – whether to set training mode (True) or evaluationmode (False). Default:True.
Returns: self
Return type: Module

type(dst_type)[source]#

Casts all parameters and buffers todst_type.

Note

This method modifies the module in-place.

Parameters: dst_type (type orstring) – the desired type
Returns: self
Return type: Module

xpu(device=None)[source]#

Move all model parameters and buffers to the XPU.

This also makes associated parameters and buffers different objects. Soit should be called before constructing optimizer if the module willlive on XPU while being optimized.

Note

This method modifies the module in-place.

Parameters: device (int,optional) – if specified, all parameters will becopied to that device
Returns: self
Return type: Module

zero_grad(set_to_none=True)[source]#

Reset gradients of all model parameters.

See similar function undertorch.optim.Optimizer for more context.

Parameters: set_to_none (bool) – instead of setting to zero, set the grads to None.Seetorch.optim.Optimizer.zero_grad() for details.

On this page

Show Source

PyTorch Libraries

Movatterモバイル変換

Module #

Docs

Tutorials

Resources

Movatterモバイル変換

Module#

Docs

Tutorials

Resources

Module #