Работа с Module в EvoX

Работа с Module в EvoX

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

В этом руководстве представлен базовый модуль в EvoX: ModuleBase.

Введение в Module

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

Инициализация Algorithm и Problem — Настройка Monitor — Инициализация Workflow — Запуск Workflow

Этот процесс требует четырех базовых классов в 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] и которое можно изменять на месте (in-place).

Роль Module

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

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

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

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

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

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

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

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

Использование Module

В частности, существуют некоторые правила использования 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 для явного определения потока управления.