oneflow.nn.Module

Module class for building neural networks

class oneflow.nn.Module

add_module(name: str, module: Optional[oneflow.nn.module.Module]) → None

Adds a child module to the current module.

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

Parameters

• name (string) – name of the child module. The child module can be accessed from this module using the given name

• module (Module) – child module to be added to the module.

apply(fn: Callable[[oneflow.nn.module.Module], None]) → T

buffers(recurse: bool = True) → Iterator[oneflow._oneflow_internal.Tensor]

children() → Iterator[oneflow.nn.module.Module]

property consistent

cpu() → T

Moves all model parameters and buffers to the CPU.

Note

This method modifies the module in-place.

Returns

self

Return type

Module

cuda(device: Optional[Union[int, oneflow._oneflow_internal.device]] = None) → T

Moves all model parameters and buffers to the GPU.

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

Note

This method modifies the module in-place.

Parameters

device (int, optional) – if specified, all parameters will be copied to that device

Returns

self

Return type

Module

double() → T

Casts all floating point parameters and buffers to double datatype.

Note

This method modifies the module in-place.

Returns

self

Return type

Module

eval() → T

extra_repr() → str

Set the extra representation of the module

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

float() → T

Casts all floating point parameters and buffers to float datatype.

Note

This method modifies the module in-place.

Returns

self

Return type

Module

load_state_dict(state_dict: Dict[str, oneflow._oneflow_internal.Tensor], strict: bool = True)

modules() → Iterator[oneflow.nn.module.Module]

named_buffers(prefix: str = '', recurse: bool = True) → Iterator[Tuple[str, oneflow._oneflow_internal.Tensor]]

named_children() → Iterator[Tuple[str, oneflow.nn.module.Module]]

named_modules(memo: Optional[Set[oneflow.nn.module.Module]] = None, prefix: str = '')

named_parameters(prefix: str = '', recurse: bool = True) → Iterator[Tuple[str, oneflow._oneflow_internal.Tensor]]

parameters(recurse: bool = True) → Iterator[oneflow._oneflow_internal.nn.Parameter]

register_buffer(name: str, tensor: Optional[oneflow._oneflow_internal.Tensor], persistent: bool = True) → None

register_forward_hook(hook: Callable[[…], None]) → None

register_forward_pre_hook(hook: Callable[[…], None]) → None

register_parameter(name: str, param: Optional[oneflow._oneflow_internal.nn.Parameter]) → None

state_dict(destination=None, prefix='', keep_vars=False) → Dict[str, oneflow._oneflow_internal.Tensor]

to(device: Optional[Union[str, oneflow._oneflow_internal.device]] = None)

to_consistent(placement=None, sbp=None)

train(mode: bool = True) → T