透明检查点实现自动微分
1. Checkpointing.jl 的设计目标
Checkpointing.jl 的目标是为时间步循环中的伴随检查点提供一个完全透明且灵活的解决方案,涵盖以下几个方面:
- 自动存储和恢复检查点变量。
- 支持多种检查点方案。
- 支持各种类型的存储设备。
2. 模型对象
为了实现上述目标,在 Julia 中定义了时间步模型的标准结构。代码使用上下文模型对象来存储当前模型的状态,这种编写代码的风格在 Julia 中很常见,因为它允许利用语言的多重派发特性对对象类型进行方法派发。以 Burgers 模型为例,其结构定义如下:
mutable struct Burgers
nextu::Matrix{Float64}
nextv::Matrix{Float64}
lastu::Matrix{Float64}
lastv::Matrix{Float64}
Nx::Int
Ny::Int
µ::Float64
dx::Float64
dy::Float64
dt::Float64
tsteps::Int
...
end
模型包含 u 和 v 的矩阵类型字段,以及每个字段的下一个和最后一个存储位置,用于模板计算(下一个值由最后一个值计算得出)。此外,还包含所有模型参数 ν、dt、dx、dy 以及网格大小 Nx 和 Ny。虽然只有 u 和 v 字段需要检查点,但为了简化操作,对整个结构体进行检查点操作,这样可以实现伴随检