Skip to content

treex.Treex

A Tree class with all Mixin Extensions. Base class for all Treex classes.

Source code in treex/treex.py
class Treex(to.Tree, to.Extensions):
    """
    A Tree class with all Mixin Extensions. Base class for all Treex classes.
    """

    def train(self: T, mode: bool = True, inplace: bool = False) -> T:
        """
        Creates a new module with the same structure, but with `Module.training` set to the given value.

        Arguments:
            mode: The new training mode.
            inplace: Whether to update the module inplace.
        Returns:
            The new module in with the training mode is set to the given value,
            if `inplace` is `True` then `self` is returned.
        """

        def set_training(tree):
            if isinstance(tree, Treex) and hasattr(tree, "_training"):
                tree._training = mode

        return to.apply(set_training, self, inplace=inplace)

    def eval(self: T, inplace: bool = False) -> T:
        """
        Creates a new module with the training mode set to False, equivalent to calling `train(False)`.

        Returns:
            The new module with the training mode set to False.
        """
        return self.train(False, inplace=inplace)

    def freeze(self: T, mode: bool = True, inplace: bool = False) -> T:
        """
        Creates a new module with the same structure, but with `Module.frozen` set
        to the given value.

        Arguments:
            mode: The new `frozen` mode.
            inplace: Whether to update the module inplace.
        Returns:
            The new module in with the `frozen` mode is set to the given value,
            if `inplace` is `True` then `self` is returned.
        """

        def set_frozen(tree):
            if isinstance(tree, Treex) and hasattr(tree, "_frozen"):
                tree._frozen = mode

        return to.apply(set_frozen, self, inplace=inplace)

    def unfreeze(self: T, inplace: bool = False) -> T:
        """
        Creates a new module with `.frozen` set to False, equivalent to calling `freeze(False)`.

        Arguments:
            inplace: Whether to update the module inplace.
        Returns:
            The new module with `.frozen` set to False, if `inplace` is `True` then `self` is returned.
        """
        return self.freeze(False, inplace=inplace)

eval(self, inplace=False)

Creates a new module with the training mode set to False, equivalent to calling train(False).

Returns:

Type Description
~T

The new module with the training mode set to False.

Source code in treex/treex.py
def eval(self: T, inplace: bool = False) -> T:
    """
    Creates a new module with the training mode set to False, equivalent to calling `train(False)`.

    Returns:
        The new module with the training mode set to False.
    """
    return self.train(False, inplace=inplace)

freeze(self, mode=True, inplace=False)

Creates a new module with the same structure, but with Module.frozen set to the given value.

Parameters:

Name Type Description Default
mode bool

The new frozen mode.

True
inplace bool

Whether to update the module inplace.

False

Returns:

Type Description
~T

The new module in with the frozen mode is set to the given value, if inplace is True then self is returned.

Source code in treex/treex.py
def freeze(self: T, mode: bool = True, inplace: bool = False) -> T:
    """
    Creates a new module with the same structure, but with `Module.frozen` set
    to the given value.

    Arguments:
        mode: The new `frozen` mode.
        inplace: Whether to update the module inplace.
    Returns:
        The new module in with the `frozen` mode is set to the given value,
        if `inplace` is `True` then `self` is returned.
    """

    def set_frozen(tree):
        if isinstance(tree, Treex) and hasattr(tree, "_frozen"):
            tree._frozen = mode

    return to.apply(set_frozen, self, inplace=inplace)

train(self, mode=True, inplace=False)

Creates a new module with the same structure, but with Module.training set to the given value.

Parameters:

Name Type Description Default
mode bool

The new training mode.

True
inplace bool

Whether to update the module inplace.

False

Returns:

Type Description
~T

The new module in with the training mode is set to the given value, if inplace is True then self is returned.

Source code in treex/treex.py
def train(self: T, mode: bool = True, inplace: bool = False) -> T:
    """
    Creates a new module with the same structure, but with `Module.training` set to the given value.

    Arguments:
        mode: The new training mode.
        inplace: Whether to update the module inplace.
    Returns:
        The new module in with the training mode is set to the given value,
        if `inplace` is `True` then `self` is returned.
    """

    def set_training(tree):
        if isinstance(tree, Treex) and hasattr(tree, "_training"):
            tree._training = mode

    return to.apply(set_training, self, inplace=inplace)

unfreeze(self, inplace=False)

Creates a new module with .frozen set to False, equivalent to calling freeze(False).

Parameters:

Name Type Description Default
inplace bool

Whether to update the module inplace.

False

Returns:

Type Description
~T

The new module with .frozen set to False, if inplace is True then self is returned.

Source code in treex/treex.py
def unfreeze(self: T, inplace: bool = False) -> T:
    """
    Creates a new module with `.frozen` set to False, equivalent to calling `freeze(False)`.

    Arguments:
        inplace: Whether to update the module inplace.
    Returns:
        The new module with `.frozen` set to False, if `inplace` is `True` then `self` is returned.
    """
    return self.freeze(False, inplace=inplace)