net

sorix.nn.net

Module

Module()

Base class for all neural network modules.

Your models should also subclass this class.

Source code in sorix/nn/net.py

def __init__(self) -> None:
    super().__init__()
    self.device: str = 'cpu'
    self.training: bool = True

xp property

xp

Returns the appropriate array module (numpy or cupy) for the module's device.

forward

forward(x)

Defines the computation performed at every call. Should be overridden by all subclasses.

Source code in sorix/nn/net.py

def forward(self, x: Tensor) -> Tensor:
    """
    Defines the computation performed at every call.
    Should be overridden by all subclasses.
    """
    raise NotImplementedError("You must implement forward in the subclass.")

parameters

parameters()

Returns an iterator over module parameters (tensors that require gradients).

Source code in sorix/nn/net.py

def parameters(self) -> List[Tensor]:
    """
    Returns an iterator over module parameters (tensors that require gradients).
    """
    params: List[Tensor] = []
    visited: Set[int] = set()

    def _gather_params(obj: Any) -> None:
        if id(obj) in visited:
            return
        visited.add(id(obj))

        if isinstance(obj, Tensor):
            if obj.requires_grad:
                params.append(obj)
        elif hasattr(obj, "parameters") and callable(obj.parameters) and obj is not self:
            # If the object has its own parameters() method, use it
            params.extend(obj.parameters())
        elif hasattr(obj, "__dict__"):
            # Recurse into attributes
            for k, v in obj.__dict__.items():
                if not k.startswith('_'):
                    _gather_params(v)
        elif isinstance(obj, (list, tuple)):
            for item in obj:
                _gather_params(item)
        elif isinstance(obj, dict):
            for item in obj.values():
                _gather_params(item)

    _gather_params(self)
    return params

to

to(device)

Moves all model parameters and buffers to the specified device.

Parameters:

• device (str) –

  'cpu' or 'cuda'.

Source code in sorix/nn/net.py

def to(self, device: str) -> Module:
    """
    Moves all model parameters and buffers to the specified device.

    Args:
        device: 'cpu' or 'cuda'.
    """
    self.device = device

    def _apply(obj: Any) -> Any:
        if hasattr(obj, "to") and callable(obj.to) and obj is not self:
            return obj.to(device)

        if isinstance(obj, list):
            return [_apply(v) for v in obj]
        if isinstance(obj, tuple):
            return tuple(_apply(v) for v in obj)
        if isinstance(obj, dict):
            return {k: _apply(v) for k, v in obj.items()}

        if hasattr(obj, "__dict__") and obj is not self:
             # Try to move attributes of non-Module objects
             for k, v in obj.__dict__.items():
                 if not k.startswith('_'):
                     setattr(obj, k, _apply(v))

        return obj

    for k, v in self.__dict__.items():
        if not k.startswith('_'):
            setattr(self, k, _apply(v))

    return self

train

train()

Sets the module in training mode.

Source code in sorix/nn/net.py

def train(self) -> None:
    """Sets the module in training mode."""
    self.training = True
    def _apply(obj: Any) -> None:
        if hasattr(obj, "train") and callable(obj.train) and obj is not self:
            obj.train()
        elif hasattr(obj, "training"):
            obj.training = True

        if isinstance(obj, (list, tuple)):
            for o in obj: _apply(o)
        elif isinstance(obj, dict):
            for v in obj.values(): _apply(v)

    for k, v in self.__dict__.items():
        if not k.startswith('_'):
            _apply(v)

eval

eval()

Sets the module in evaluation mode.

Source code in sorix/nn/net.py

def eval(self) -> None:
    """Sets the module in evaluation mode."""
    self.training = False
    def _apply(obj: Any) -> None:
        if hasattr(obj, "eval") and callable(obj.eval) and obj is not self:
            obj.eval()
        elif hasattr(obj, "training"):
            obj.training = False

        if isinstance(obj, (list, tuple)):
            for o in obj: _apply(o)
        elif isinstance(obj, dict):
            for v in obj.values(): _apply(v)

    for k, v in self.__dict__.items():
        if not k.startswith('_'):
            _apply(v)

state_dict

state_dict()

Returns a dictionary containing the whole state of the module (parameters and buffers).

Source code in sorix/nn/net.py

def state_dict(self) -> Dict[str, Tensor]:
    """
    Returns a dictionary containing the whole state of the module (parameters and buffers).
    """
    state: Dict[str, Tensor] = {}

    def _get_state(obj: Any, prefix: str) -> None:
        for name, val in obj.__dict__.items():
            if name.startswith('_') or name == 'device' or name == 'training':
                continue

            key = prefix + name
            if isinstance(val, Tensor):
                state[key] = val
            elif hasattr(val, 'state_dict') and callable(val.state_dict) and val is not self:
                sub_state = val.state_dict()
                for sk, sv in sub_state.items():
                    state[key + '.' + sk] = sv
            elif hasattr(val, '__dict__'):
                 for kn, kv in val.__dict__.items():
                     if isinstance(kv, Tensor):
                         state[key + '.' + kn] = kv

    _get_state(self, "")
    return state

extra_repr

extra_repr()

Set the extra representation of the module. Should be overridden by subclasses.

Source code in sorix/nn/net.py

def extra_repr(self) -> str:
    """
    Set the extra representation of the module.
    Should be overridden by subclasses.
    """
    return ""

load_state_dict

load_state_dict(state_dict)

Copies parameters and buffers from state_dict into this module and its descendants.

Source code in sorix/nn/net.py

def load_state_dict(self, state_dict: Dict[str, Tensor]) -> None:
    """
    Copies parameters and buffers from state_dict into this module and its descendants.
    """
    own_state = self.state_dict()
    for name, param in state_dict.items():
        if name in own_state:
            if isinstance(param, Tensor):
                own_state[name].data = param.data
                own_state[name].to(own_state[name].device) # Ensure device consistency
            else:
                pass
        else:
            pass

Sequential

Sequential(*args)

Bases: Module

A sequential container. Modules will be added to it in the order they are passed in the constructor.

Examples:

model = nn.Sequential(
    nn.Linear(10, 5),
    nn.ReLU(),
    nn.Linear(5, 2)
)

Source code in sorix/nn/net.py

def __init__(self, *args: Any) -> None:
    super().__init__()
    self._modules: Dict[str, Module] = {}
    if len(args) == 1 and isinstance(args[0], dict):
        for name, module in args[0].items():
            setattr(self, name, module)
        self._modules = args[0]
    else:
        for idx, module in enumerate(args):
            name = str(idx)
            setattr(self, name, module)
            self._modules[name] = module

ModuleList

ModuleList(modules=None)

Bases: Module

Holds submodules in a list. ModuleList can be indexed like a regular Python list, but modules it contains are properly registered, and will be visible by all Module methods.

Examples:

self.layers = nn.ModuleList([nn.Linear(10, 10) for i in range(5)])

Source code in sorix/nn/net.py

def __init__(self, modules: Optional[List[Module]] = None) -> None:
    super().__init__()
    self._items: List[Module] = []
    if modules is not None:
        self.extend(modules)