backend 模块详解

backend 模块详解#

后端模块独占 FEALPy 软件架构中的张量层（见下图），向下适配多种张量计算框架，向上支撑 FEALPy 中所有的张量化算法。目前支持的张量计算后端有 NumPy、PyTorch、JAX。

为什么我们要支持多种计算后端？每个领域都有其倾向的、惯用的张量计算库，每种计算库都对不同厂商、平台的硬件做适配和优化。因此，支持更多的张量计算库就是适配更多计算硬件、对接更多计算领域。

后端模块的核心功能是切换 FEALPy 中张量计算工作所用的计算库，即对于 FEALPy 而言的“计算后端”。基于此目标，后端模块被设计得十分易用。

后端管理器的使用#

从后端模块里导入后端管理器，即可完全控制 FEALPy 的后端模块：

from fealpy.backend import backend_manager as bm

最基本地，FEALPy 默认的张量计算后端是 NumPy，利用 set_backend 函数可调换所用后端：

bm.set_backend('pytorch')

后端管理器自身按 Python Array Standard 注册有大量张量计算函数，可供用户直接调用，执行当前后端下的张量计算，且在使用习惯上与 NumPy、PyTorch 等计算库一致，例如：

bm.set_backend('pytorch')
bm.zeros([1, 2, 3], dtype=bm.float32)
# 此时相当于 torch.zeros([1, 2, 3], dtype=bm.float32)

另外，bm 对后端的调换操作对所有 FEALPy 上层张量相关模块生效。示例：对比不同后端下构造的网格对象。

① 设置后端为 NumPy，然后构造均匀的三角形网格，查看网格顶点和单元的数据类型。#

from fealpy.mesh import TriangleMesh
bm.set_backend('numpy')

mesh = TriangleMesh.from_box()
print(type(mesh.entity('cell')))
print(type(mesh.entity('node')))
# output:
# <class 'numpy.ndarray'>
# <class 'numpy.ndarray'>

② 设置后端为 PyTorch，然后构造均匀的三角形网格，查看网格顶点和单元的数据类型。#

from fealpy.mesh import TriangleMesh
bm.set_backend('pytorch')

mesh = TriangleMesh.from_box()
print(type(mesh.entity('cell')))
print(type(mesh.entity('node')))
# output:
# <class 'torch.Tensor'>
# <class 'torch.Tensor'>

Backend 模块是 FEALPy 的底层控制，能在不改变用户脚本的情况下更换张量计算库。

附录：后端模块架构#

alt text

  • 管理器：负责载入后端、切换后端，将计算指令传递给当前选定的适配器。

  • 适配器：接收满足标准接口的计算指令，转换成对应后端的函数调用。

Contents