基于Transformer 结构的高精度湍流波前重构

冯佳濠，胡启立，姜 律，杨燕燕，华晟骁，吴晶晶，胡立发*

（1.江南大学 理学院，江苏 无锡 214122；2.江苏省轻工光电工程技术研究中心，江苏 无锡 214122；3.光电对抗测试评估技术重点实验室，河南 洛阳 471003）

1 引 言

在大口径光学望远镜的天文观测中，大气湍流等观测条件变化严重限制了望远镜的成像分辨率，因此人们引入自适应光学（AO）技术来克服这一问题，以尽可能接近其衍射极限的成像性能［1］。

自适应光学系统［2］主要由3 部分组成：波前传感器、波前控制器和波前校正器。波前传感器检测波前像差，直接决定了系统的波前校正精度。在几种典型的波前探测器中，夏克-哈特曼波前传感器（SHWFS）［3］以其高光能量利用率高、结构简单、探测动态范围大而成为应用最多的波前探测器。SHWFS 由子透镜阵列和CCD 组成。每个子透镜的波前像差使光斑的质心发生位移。根据这些质心偏移量，可以计算出相应子区域的平均斜率，从而以斜率与相位关系重构波前。

为应对强大气湍流的影响，大口径光学望远镜（如VLT 巡天望远镜［4］）上的SHWFS 需要大量的子透镜来分割湍流引起的畸变波前，提高其探测精度，但由于每个子透镜光通量降低，导致其探测能力降低。此外，大量的重心（COG）［5］和斜率计算［6］会降低探测波前的重构速度且基于斜率的传统波前重构算法重构精度始终有限，因此，研究一种能够保证SHWFS 波前探测精度和波前探测能力并且能实现高精度且稳定波前重构的算法是很有必要的。

近两年，Gao 等人［7］利用基于维纳滤波器的迭代计算策略重构出在强湍流影响下的被观测物体的清晰图像。Kee 等人［8］使用机器学习的方式辅助以实现在强湍流的情况下进行数字识别。但是这两种方法并不适用于大口径天文望远镜上的SHWFS。Dubose 等人［9］扩展了Swanson 等人［10］的工作，构建了强度/斜率网络（ISNet），它可以由斜率和子透镜采集的光点的总强度在强湍流情况下直接重构波前，该方法适用于SHWFS，然而其计算斜率的过程会降低重构速度并且其波前畸变波动的剧烈程度有限。胡乐佳等人［11］减少了斜率的计算过程，构建的SH-Net 结构是已知的重构精度最高的方法，却只是在生物显微的像差测量应用中实现了从传感器的光点阵图像到相位分布的高精度映射。

上述针对强湍流影响的研究确实提高了波前重构精度或者直接提高了光学系统的成像能力，然而均没有考虑到观测目标亮度减弱的影响。李自强［12］等人构造了名为SHNN 的全连接神经网络，解决了光强干扰的极端情况下导致的质心探测精度不高或无法探测的问题，并通过有限数目的Zernike 模式重构波前，然而却存在截断误差和模式混淆误差等问题，波前重构精度不高。在天文观测中，湍流强度的变化会导致波前的波动范围很大，随机的变化也会导致波前有较多高阶像差，Zernike 模式重构波前的方法平滑了边缘高阶像差并且存在截断误差和模式混淆误差，波前重构精度有限。

为了在湍流强度强弱变化与观测目标亮度变化的复杂观测条件下，基于固定的子透镜阵列的SHWFS 实现高精度波前重构，本文提出了一个基于Transformer的U 形神经网络模型。在数值验证部分，通过模拟由大气湍流相干长度r0［13］从5～25 cm 的大范围湍流空间强度变化引起的复杂畸变波前（涵盖了世界主要天文台的典型地观测条件），对所提出的模型进行了测试和分析。通过模拟不同星等的观测条件，进一步测试了模型对于不同亮度的目标的波前重构能力。

2 波前数据和SH-U-Transformer模型

2.1 数据模拟

模拟的大气湍流［14］为3 个离散层叠加形成的湍流相位屏，湍流相位屏遵循von Kármán 谱［15］，其具体参数如表1 所示。von Kármán 谱的表达式如式（1）所示：

表1 湍流相位屏的关键参数Tab.1 Key parameters for turbulence phase screens

其中：f为空间频率；L0为湍流外尺度；r0为大气相干长度［16］，即湍流空间强度。通过Bufton 模型［17］来设置风速参数，其定义如式（2）所示：

其中h为高度，其单位为m。

为了使训练后的模型在变化的观测条件下能够获得较高精度的畸变波前重构，在湍流模拟过程中，将湍流空间强度最强的情况，即大气相干长度r0的最小值设置在V 波段（505～595 nm）的5 cm 处，以适应一般天文台需要面对的最强的大气湍流扰动。导星所处波段也为V 波段（激光或瑞利导星所处波段）并将SHWFS 设置为V 波段。观测目标的的星等为6 mag。

模拟的望远镜口径为2 m，在模拟的SHWFS［18］上建立了一个大规模的子透镜阵列。子透镜阵列由40×40（约1 208 个有效子透镜）子透镜组成，光点阵图像总像素数为280×280 个，单个像素直径为10 μm。在模拟CCD 采集光点阵图像的过程中，考虑了3 种主要噪声类型；暗电流噪声、读出噪声和光子噪声。

2.2 SH-U-Transformer 结构

在复杂多变的大气湍流环境下利用光点阵图像实现复杂畸变波阵面的直接高精度重构，一般使用U 型网络结构且需要网络具有极强表达能力。然而已报道的应用于波前重构的CNN，包括SH-Net在内，由于局部卷积性质，使得其每层仅能关注到卷积核内的部分，光点之间的联系往往无法被网络学习，除非网络深度极深，这会导致训练过程出现信息丢失，过于庞大的网络结构需要更高的计算和时间成本，并不适用于波前重构。

本文为了网络能在复杂观测条件下建立足够稳定的非线性映射关系，需要其对输入的光点阵图像在关注局部光点有效信息的同时，进行全局所有光点信息和噪声的交互式学习以实现从单一光点阵中获得足够丰富的特征信息。而网路输出的畸变波前整体是复杂且具有较多的高频信号，并且随着观测条件的变化而变化，这就需要网络对输入的特征图具有强大的非线性解码能力，才能够实现对畸变波前的高准确度重构。

Transformer 结构［19］以自注意力机制［20］为中心，能够对局部细节进行有效关注并使全局信息实现交互。由于这种全局建模的能力，本文首次将Transformer 结构引入到波前重构中。选择的Transformer模块为Swin-Transformer模块［21］，相比于普通的Transformer 模块，Swin-Transformer 模块基于移动窗口的自注意力机制使其不仅能够进行全局建模，而且计算复杂度更低，有助于网络在重构波前时拥有更快的响应速度。

本文构建了编码器和解码器均以Transformer 模块为主学习模块的U 型网络结构，并将其命名为SH-U-Transformer，其结构如图1（a）所示。SH-U-Transformer 网络模型的输入为SHWFS 中280×280 大小的光点阵图像，输出为280×280 大小的重构波前。

图1 SH-U-Transformer 的架构和细节。（a）SH-U-Transformer 的架构（H 和W 表示特征图的大小，C 是通道数）；（b）Embedding Layer 的结构（k 是卷积核的大小，s 是卷积的步长，p 是填充大小）；（c）Transformer Block 的结构（多头自注意力机制中的注意力头数随着通道数的增加而增加1 倍，初始注意头数为9）。Fig.1 Architecture and details of SH-U-Transformer.（a）Architecture of SH-U-Transformer（H and W indicate the size of the feature map.C is the channel’s number）；（b）Architecture of the Embedding Layer（k is the size of the convolution kernel，s is the stride of the convolution，and p is the padding added to all four sides of the input）；（c）Architecture of the Transformer Block（the number of attentional heads in the two multi-headed self-attention mechanisms doubles as the number of channels increases，and the initial number of attention heads is 9）.

为了确保没有信息丢失，该网络模型的特征图尺寸调整部分并没有选择类似于CNN 粗暴的信息筛选或填充的方式。在编码器部分，模型通过分割相邻像素块然后将它们通过前后通道拼接重排成新的特征块来进行下采样，如图2 所示；在解码器部分相反，模型将相邻通道的特征块重新左右拼接使通道数量减半，特征图的分辨率变为原来的2 倍。这一做法确保了所有信息都能被运用并通过学习的方式决定去留。

图2 SH-U-Transformer 下采样的方式Fig.2 SH-U-Transformer downsampling mode

接下来是对模型进行进一步优化调整。首先是编码器部分的预处理层——Embedding Layer，这一层的作用是对光点阵图像进行初步特征提取和通道增加的预处理后输出的特征图进行特征学习提取。值得一提的是，在该层用于初步特征提取的部分，本文自行设计了一个Inception Block［22］，如图1（b）所示，用于对输入进行预处理。在这个Block 中，设计有4 种不同的卷积内核，它们的尺寸分别是5×5、7×7、9×9 和11×11。通过多感受野融合，该模块实现了从提取单个子透镜区域的光点信息扩展到与相邻子透镜区域的光点信息交互或者单个子透镜内的有效光点和噪声信息交互。通过将这些提取出来的信息通过通道拼接的方式输入到网络中可以保证输入网络的特征图携带更多的信息供给编码器的学习模块进行学习。

其次，对图1（c）中每个Transformer Block 的学习核心——多头自注意力机制中的注意力头的数量进行修改。多头自注意力机制中的头能够在聚焦到像素点自身的同时关注周围区域的像素点，这也是Transformer Block 在建立每个子透镜中光点与子透镜区域内波前的像素级映射的同时学习与其他区域波前之间的关联的核心。在一个自注意力机制中，其头数需要大于等于9 才能使每个机制学习能力均超过卷积模块［23］，而一个Transformer Block 拥有两个类似的自注意力机制，这确保了每个Transformer Block 的学习能力均远超卷积模块。本文在第一个包含Transformer Block 的学习层将注意力头数设置为9，并且随着网络的深入，注意力头数倍增，每一层的学习能力增强，从而进一步提高了整个网络的表达能力。

最后是整个网络的Transformer Block 的布局。所有的Transformer Block 具有相同的结构，并且是模型实现高精度直接重构的关键。编码器和解码器各设置了一个叠加6 个Transformer Block 的学习层。编码器部分被设置为第二个主要的学习层，用于从下采样过一次的特征图中学习特征，这是因为经过初步特征提取和首次学习的特征图包含了足够多的光点阵特征与联系，足够多的学习模块有助于对这些信息进行辨析提取和进一步学习，防止网络加深导致信息丢失。解码器部分被设置为倒数第二个主要学习层，是从特征图逐步解码为波前分布的关键部分。输入这一层的特征图包含了大量与波前分布相关的信息，需要网络的进一步解析后由最后一个学习层进行调整进而输出波前。Transformer Block 数量的增加可以增强每个学习层的学习能力，然而也会增加参数量。为了在减少时间成本的同时保证精度，其他的学习层的Transformer Block 的数量经过测试训练后都被设置为2 个。

2.3 神经网络的训练

训练数据集包括90 000 组对应的SHWFS 的光点阵图像及畸变波前。其中的90%作为训练数据集，用于执行算法训练，优化SH-U-Transformer模型的参数；10%的数据作为验证数据集，用于充分验证以选择最佳模型。SH-U-Transformer 网络模型使用Pytorch 框架（Python 3.8）构建，在桌面PC（3.20 GHz Intel（R）Core i7-8700 CPU，16 GB RAM 和NVIDIA GeForce RTX 1070）上进行训练。

在训练过程中，由于Transformer Block 的计算复杂度较高，SH-U-Transformer 模型需要更多的训练批次和更详细且有针对性的损失函数设置，因此需要将整个训练过程分为训练和调整两个阶段，两个阶段的批次大小均为6 个。为了使训练后的模型在重构高度复杂的波前时更加精确、稳定，损失函数被设置为复合函数：

其中：t为真实波前；p为网络重构的波前；RMSE为均方根误差函数；SSIM［24］是结构相似性函数，其作为权重被添加到损失函数中；α是一个超参数，它决定了结构相似性函数占损失函数的比例。在第一阶段设置有200 个epoch，初始学习率设置为1×10-4，α设置为实验测量的0.013。采用Adam 优化器进行快速梯度下降，快速优化模型参数。在第二阶段设置有50 个epoch，初始学习率设置为6×10-7，α设置为0.003 5。优化器为SGD 优化器，对模型参数进一步调整优化。在这两个阶段中，模型验证阶段的评估指标均为真实波前和预测波前之间的RMSE 值作为最终的验证准确度，RMSE=0.010 77 μm。

3 数值模拟结果与讨论

3.1 不同湍流空间强度对波前重构精度的影响

训练好的SH-U-Transformer 模型参数被保存后用于波前重构，仿真光路如图3 所示。在模拟设置中，来自导星或探测目标的光波通过地球大气湍流后由望远镜的镜面收集。望远镜收集的光波经过缩束处理被反射镜反射后被SHWFS采集。在此过程中，大气湍流扰动的变化引起波前像差的变化以及导星或探测目标星等的升高对SHWFS 探测能力的影响，这两点都降低了望远镜的观测能力。

图3 用于数据生成和波前重建的模拟光路图Fig.3 Simulated optical path diagram for data generation and wavefront reconstruction

在SH-U-Transformer 模型性能测试阶段，通过随机生成若干大气相干长度r0处于5～20 cm 范围内的SHWFS 接收光点阵图像作为测试数据，输入训练好的网络模型，测试在动态变化湍流空间强度下的模型重构性能。首先选择使用两种常规的波前重建方法：区域法［25］和前250 项Zernike 模式的模态法［26］进行比较。前250项Zernike 模式的系数分布如表2所示。另外，本文复刻了SH-Net并将其以训练数据进行训练，训练好后将其与SHU-Transformer 模型进行实验对比。

表2 前250 项Zernike 多项式系数分布Tab.2 The first 250 terms Zernike polynomial coefficient distribution

在图4 中，从上到下，每一行对应一个有代表性的r0条件，分别为5，10，15，20 cm。对应条件下的光点阵图像即是网络的输入，输出是相应的畸变波前。第一列为光点阵图像，第二列为对应的真实波前，第三列为SH-U-Transformer 模型重构的波前，第四列为该方法对应的残余波前。后面3 组从前至后依次为SH-Net 方法、模态法和区域法的结果和残差。可以看出，在不同的r0条件下，SH-U-Transformer 模型波前重构精度的峰谷值（Peak-to-Valley，PV）与均方根误差（RMSE）都是最优的。例如，当湍流空间强度在r0=5 cm 时，SH-U-Transformer 模型计算的残余波前的PV 值为0.265 1 μm，约为真实波前PV 值（19.788 4 μm）的1.34%；RMSE 值为0.023 9 μm，约为真实波前RMSE 值（4.123 3 μm）的0.58%。与SH-Net 相比，残余波前的PV 值约比其低57.25%（0.620 1 μm），RMSE 值约比其低51.32%（0.049 1 μm）；与模态法相比，残余波前的PV 值约比其低98.46%（17.173 8 μm），RMSE 值约比其低94.42%（0.428 6 μm）；与区域法相比，残余波前的PV 值约比其低97.80%（12.075 9 μm），RMSE值约比其低97.21%（0.856 6 μm）。该方法即使是在空间强度极强的湍流下，波前重构精度也能远优于已有的方法。虽然随着湍流空间强度减弱而优势减弱，但其重构精度却一直十分稳定。

图4 4 种方法在4 种不同湍流强度下波前重构结果的比较。计算的波前及其残差使用了与原始波前不同的颜色分布标准，给出了每个残差的均方根误差（RMSE）和PV 误差。Fig.4 Comparison of the wavefront reconstruction results of the four methods at four different r0.The calculated wavefront and its residuals use a different color bar than the original wavefront.The RMS error（RMSE）and the PV error for each residual are given.

在4 个不同区间的湍流空间强度下的400 个4种波前重构方法的残余波前RMSE 值分布如图5所示，湍流空间强度分布的规律为r0每隔5 cm 为一个区间。结果表明，在r0处于5～25 cm 区间内的SH-U-Transformer 的性能均大幅优于其他方法，残余波前的RMSE 值分布范围为0.010～0.024 μm，几乎不受湍流空间强度变化的影响。这进一步补充说明了SH-U-Transformer在湍流空间强度变化下的重构波前结果稳定性，在典型r0条件下均具有显著优势。此外，SH-U-Transformer可以直接由SHWFS 采集的光点阵图像重构出高精度的波前分布，省去了许多计算步骤，简化了波前重构的流程。

图5 4 种波前重建方法在不同湍流空间强度下残余波前RMS 误差分布。（a）r0范围为5～10 cm；（b）r0范围为10 cm～15 cm；（c）r0范围为15～20 cm；（d）r0范围为20～25 cm。每个阶段屏幕包含100 个数据。Fig.5 Comparison of RMS errors of residual wavefront under different r0 for four wavefront reconstruction methods.（a）r0 ranges from 5 cm to 10 cm；（b）r0 ranges from 10 cm to 15 cm；（c）r0 ranges from 15 cm to 20 cm；（d）r0 ranges from 20 cm to 25 cm.Each stage screen contains 100 datasets.

对评价波前重构算法优劣的关键指标——重构速度［27］也进行了测试。将SH-U-Transformer模型部署到不同的GPU 上进行重构波前计算并记录时间成本。在部署了RTX A6000 显卡的服务器上，计算时间约为15 ms，这是目前最快的计算速度，表明SH-U-Transformer 模型可以无需调整直接部署到更高型号的GPU 上以尽可能地实现实时的计算。值得注意的是，当模型部署在配备了RTX 2080ti GPU 的工作站上时，波前重构所需的时间（24.87 ms）比已报道的SH-Net（40.2 ms）少了约38.13%，表明SH-U-Transformer具有更快的波前重构速度，这也为其实际应用提供了可能。

3.2 观测目标的不同亮度对波前重构精度的影响

本文还测试了SH-U-Transformer 模型在应对不同星等［28］的导星或观测目标时重构波前的精度。SHWFS 子透镜阵列接收的光子数量随着导星或观测目标的星等的变化而变化，这也导致了有效子透镜数量的变化。表3给出了当r0=5 cm时，观测目标的星等在3～8 mag 时，有效子透镜数目及子透镜接收的平均光强。

表3 不同星等下SHWFS 子透镜的参数Tab.3 Parameters of sub-apertures of SHWFS at different magnitudes

如表3 所示，星等的增大导致来自观测目标的光强减小，噪声增大，导致SHWFS［29］有效子孔径数量迅速减少。SHWFS 的微透镜阵列无法快速改变，然而有效子透镜的减少会导致光斑阵图像中包含的有效信息减少，从而降低波前重建精度，甚至由于过少的有效子透镜出现无法重构波前的情况。自适应光学系统通常只能分得望远镜接收的部分光用于波前探测，即子透镜接收到的光能量有限，这也导致了传统依靠有效子透镜重构波前的方法无法适用于高星等观测目标，限制了固定规模子透镜阵列SHWFS 的探测能力，进一步导致传统自适应光学系统只能应用于低星等的观测目标。

首先，随机生成了一个畸变波前及其对应的在不同星等的观测目标下SHWFS 采集的光点阵图像，通过SH-U-Transformer模型进行波前重建，结果如图6 所示。图6（a）是SH-U-Transformer模型对不同星等目标的波前重构的结果，其中从上到下依次为3、5、6、7、8 等5 个星等下的导星。结果表明，训练好的模型无需调整就可以在应对不同星等的目标时都能实现高精度波前重构。即使是星等为8 等的观测目标，从图6（a）中可以看出，残余波前的RMSE 值依然可以达到0.047 5 μm，PV 值为0.450 8 μm，与真实波前相比，重构精度令人满意。图6（b）展示了5 个不同星等的观测目标的残余波前RMSE 值分布，每组包含100 个在大气相干长度r0为5～20 cm 范围内的真实波前数据。它们的分布范围直观地表明了SH-UTransformer 在不同观测目标亮度高低变化的情况下依然能稳定地实现高精度波前重构且对暗弱目标也很有效。

图6 不同星等的观测目标的波前重构结果。（a）同一波前从上到下不同星等的重构结果，输出波前及其残差使用与真实波前不同的颜色分布标准并给出了每个残差的均方根误差（RMSE）和PV 误差；（b）r0处于5～20 cm 范围内的不同星等观测目标的重构波前的RMS 误差分布。每个星等包含100 个数据。Fig.6 Results of wavefront reconstruction at different magnitudes.（a）Reconstruction results from the same wavefront at different magnitudes from top to down，the output wavefronts and their residuals use a different color bar than the original wavefront，and RMS error（RMSE）and the PV error for each residual are given；（b）RMS error distributions of reconstructed wavefronts of different magnitudes are obtained when r0 is in the range of 5～20 cm.Each stage screen contains 100 datasets.

对于这一结果，主要原因是在SH-U-Transformer 模型的训练过程中，损失函数中加入了SSIM 函数这一权值，使模型能够通过观测目标亮度固定的数据集学习光点阵图像与波前分布之间在结构上的关系，即使亮度降低导致有效信号减弱，两者结构上的非线性关系依然能起到稳定地指导建立像素级映射的作用。因此，采用SH-U-Transformer 模型的SHWFS 能够适应导星或观测目标亮度变化的观测条件。

4 结 论

本文实现了一种深度学习波前重建方法——SH-U-Transformer。该方法适用于大口径天文光学望远镜，能够应对复杂的观测条件。SH-UTransformer 首次在波前重构中引入了Transformer 结构，实现了在复杂变化的观测条件下直接从SHWFS 光点阵图像高精度地重构湍流引起的畸变波前。通过仿真验证，定量评价了SHU-Transformer 对不同空间强度的湍流引起的像差的波前重构性能。与传统的模态法和区域法相比，SH-U-Transformer 模型不仅精度提高了近2 个数量级，而且消除了斜率测量等的繁琐步骤。与应用于SHWFS 的直接高精度重构波前的深度学习方法SH-Net 相比，SH-U-Transformer 模型具有更高的重构精度和更快的重构速度，并且随着GPU 的进一步升级，它的时间成本也会随之降低。此外，SH-U-Transformer 也可以稳定地高精度重构在导星或者观测目标改变时的畸变波前，性能表现极佳。