

Module warm.engine

PyWarm engine to the functional interface.

Functions

activate

def :
    x,
  spec,
  lookup=None

Activate tensors with given nonlinearity specification.

x: Tensor or list of Tensor; The tensors to be initialized.
spec: str or callable or 2-tuple; If a str, should be one of the nonlinearity functions contained in torch.nn.functional or torch. If a callable, it will be applied to x directly, i.e. spec(x). If a 2-tuple, it must be of format (callable, kwargs), i.e. callable(x, **kwargs).
lookup: None or list of module; Parent modules to look for spec. If None, [nn.functional, torch] is used.
return: Tensor or list of Tensor; Activation results.

forward

def :
    x,
  base_class,
  base_name=None,
  name=None,
  base_arg=None,
  base_kw=None,
  parent=None,
  infer_kw=None,
  in_shape='BCD',
  base_shape='BCD',
  out_shape='BCD',
  tuple_out=False,
  forward_arg=None,
  forward_kw=None,
  initialization=None,
  activation=None,
  **kw

A forward template that creates child instances at the first time it is called.

x: Tensor; The nd-tensor to be forwarded.
base_class: Module; A child torch.nn.Module that will be created at the first time this function is called.
base_name: str; Name for the base_class. Default: base_class name.
name: str; Name for the child module instance. Default: class name plus ordinal.
base_arg: tuple; Positional args to be passed to create the child module instance. Default: None.
base_kw: dict; KWargs to be passed to create the child module instance. Default: None.
parent: Module; The parent of the child instance. Default: None. If None, will use get_default_parent.
infer_kw: dict; Key should be valid for the child instance. Value shoud be a character, one of 'B', 'C', or 'D' (see permute), to substitute for a dimension of x. Default: None.
in_shape: str; The dimension shape of x. See also permute. Default: 'BCD'.
base_shape: str; The dimension shape required by the child module. See also permute. Default: 'BCD'.
out_shape: str or tuple or None; The dimension shape of returned tensor. See also permute. Default: 'BCD'.
tuple_out: bool; Whether the child module will return more than 1 outputs (e.g. nn.RNN). If True, the returned value of the function will be a tuple containing all outputs. Default: False.
forward_arg: tuple; positional args to be passed when calling the child module instance. Default: None.
forward_kw: dict; KWargs to be passed when calling the child module instance. Default: None.
initialization: dict; Keys are name of parameters to initialize. Values are init specs, which can be a, str, a callable, or 2-tuple; See the spec argument of initialize_ for details. Default: None.
activation: str or callable or 2-tuple; See the spec argument of activate. Default: None.
return: Tensor or tuple; If tuple_out is True, the returned value will be a tuple.

get_default_parent

def :

Get the default parent module.

initialize_

def :
    x,
  spec

Initialize parameters with given nonlinearity specification.

x: Tensor or list of Tensor; The tensors to be initialized.
spec: str or callable or 2-tuple; If a str, should be one of the nonlinearity functions contained in torch.nn.init. If a callable, it will be applied to x directly, i.e. spec(x). If a 2-tuple, it must be of format (callable, kwargs), i.e. callable(x, **kwargs).

is_ready

def :
    model

Check if a model is prepared.

namespace

def :
    f

After decoration, the function name and call count will be appended to the name kw.

permute

def :
    x,
  in_shape='BCD',
  out_shape='BCD',
  **kw

Permute the dimensions of a tensor.

x: Tensor; The nd-tensor to be permuted.
in_shape: str; The dimension shape of x. Can only have characters 'B' or 'C' or 'D', which stand for Batch, Channel, or extra Dimensions. The default value 'BCD' means the input tensor x should be at lest 2-d with shape (Batch, Channel, Dim0, Dim1, Dim2, ...), where Dim0, Dim1, Dim2 ... stand for any number of extra dimensions.
out_shape: str or tuple or None; The dimension shape of returned tensor. Default: 'BCD'. If a str, it is restricted to the same three characters 'B', 'C' or 'D' as the in_shape. If a tuple, in_shape is ignored, and simply x.permute(out_shape) is returned. If None, no permution will be performed.
return: Tensor; Permuted nd-tensor.

prepare_model_

def :
    model,
  *data,
  device='cpu'

Initialize all childen modules defined by warm in a parent model.

model: Module; The parent model to be prepared.
data: Tensor, or list of int; A batch of data with the correct shape and type to be forwarded by model. data can also be a list of int, in which case it is interpreted as the shape of the input data.
device: str, or torch.device; Should be the same for model and data. Default: 'cpu'.
return: Module; The prepared model, with all children modules defined by warm initialized.

set_default_parent

def :
    parent

Set the default parent module.

unused_kwargs

def :
    kw

Filter out entries used by forward and return the rest.