在 EvoX 中使用模块

一个模块是编程中的一个基本概念，指的是一个自包含的代码单元，旨在执行特定任务或一组相关任务。

此笔记本将介绍 EvoX 中的基本模块：ModuleBase。

模块介绍

在教程中，我们提到了EvoX中的基本运行流程：

Initiate an algorithm and a problem -- Set an monitor -- Initiate a workflow -- Run the workflow

在EvoX中，此过程需要四个基本类：

• Algorithm

• Problem

• Monitor

• Workflow

有必要为它们提供一个统一的模块。在EvoX中，这四个类都继承自基础模块——ModuleBase。

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] 中暴露出来，可以通过 self.[name] = [值] 来修改。

模块的作用

在 EvoX 中，ModuleBase 可以帮助：

• 包含可变值

该模块是一个面向对象的模块，可以包含可变值。

• 支持函数式编程

支持函数式编程模型通过 self.state_dict() 和 `self.load_state_dict(...)

• 标准化初始化:

基本上，预定义的子模块将被添加到此模块中，并在成员方法中访问，应该被视为“非静态成员”，而其他任何成员都应该被视为“静态成员”。

建议将非静态成员的模块初始化写在重写的 setup 方法（或任何其他成员方法）中，而不是 __init__ 中。

模块使用

具体来说，在EvoX中使用ModuleBase有一些规则：

静态方法

要将静态方法定义为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

非静态方法

如果在使用 vmap 时需要使用带有 Python 动态控制流（例如 if）的方法，请使用 torch.cond 明确定义控制流。