Работа с модулями в EvoX

Работа с модулями в EvoX

Модуль — это фундаментальная концепция в программировании, обозначающая самодостаточную единицу кода, предназначенную для выполнения конкретной задачи или набора связанных задач.

Этот ноутбук познакомит вас с базовым модулем в EvoX: ModuleBase.

Введение в модули

В учебном пособии мы упоминали базовый процесс работы в EvoX:

Инициализация алгоритма и задачи -- Настройка монитора -- Инициализация рабочего процесса -- Запуск рабочего процесса

Этот процесс требует четырёх базовых классов в EvoX:

  • Algorithm
  • Problem
  • Monitor
  • Workflow

Необходимо предоставить единый модуль для них. В EvoX все четыре класса наследуются от базового модуля — ModuleBase.

Module base

Класс ModuleBase

Класс ModuleBase наследуется от torch.nn.Module.

В этом классе много методов, и вот некоторые важные:

МетодСигнатураИспользование
__init__(self, ...)Инициализация модуля.
load_state_dict(self, state_dict: Mapping[str, torch.Tensor], copy: bool = False, ...)Копирование параметров и буферов из state_dict в этот модуль и его потомков. Переопределяет torch.nn.Module.load_state_dict.
add_mutable(self, name: str, value: Union[torch.Tensor | nn.Module, Sequence[torch.Tensor | nn.Module], Dict[str, torch.Tensor | nn.Module]]) -> NoneОпределение изменяемого значения в этом модуле, доступного через self.[name] и модифицируемого на месте.

Роль модуля

В EvoX ModuleBase помогает:

  • Содержать изменяемые значения

    Этот модуль является объектно-ориентированным и может содержать изменяемые значения.

  • Поддерживать функциональное программирование

    Модель функционального программирования поддерживается через self.state_dict() и self.load_state_dict(...).

  • Стандартизировать инициализацию:

    В основном, предопределённые подмодули, которые будут ДОБАВЛЕНЫ к этому модулю и доступны позже в методах-членах, должны рассматриваться как «нестатические члены», тогда как любые другие члены должны рассматриваться как «статические члены».

    Инициализацию модуля для нестатических членов рекомендуется писать в переопределённом методе setup (или любом другом методе-члене), а не в __init__.

Использование модуля

В частности, есть некоторые правила использования ModuleBase в EvoX:

Статические методы

Статические методы для JIT должны быть определены следующим образом:

# One example of the static method defined in a Module

@jit
def func(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    return x + y

Нестатические методы

Если метод с динамическим управлением потоком Python, таким как if, должен использоваться с vmap, пожалуйста, используйте torch.cond для явного определения управления потоком.