论文标题:ABC-Net: a divide-and-conquer based deep learning architecture for SMILES recognition from molecular images
代码:https://github.com/zhang-xuan1314/ABC-Net/
现有的OCSR方法由于其恢复精度较差,远远低于现实要求的期望。为能产生SMILES,模型需要识别原子和键,并以特定的顺序排列原子。这极大地增加了识别难度,导致这种类型的模型的数据效率非常低,需要百万级的训练数据来达到良好的性能。此外,随着SMILES长度的增加,模型性能会迅速恶化。
数据集:ChEMBL database,使用RDKit实现SMILES筛选。分子中超过50个非氢原子被剔除。随机抽取10万种化合物。
利用分治原则,将识别分子的问题转化为识别原子和键的问题:(A)原子和键的检测、(B)原子和键的辅助信息识别、(C)分子的重建
在模型训练阶段,使用多任务训练方法同时训练原子和键的检测任务以及其他辅助属性识别任务
一、模型:
将预测目标热图中的所有正像素点设为1,所有正像素点的一阶邻居设为0.95,其他位置设为0。当两个键中心在热图中太近甚至重叠时,这可能会导致键缺失,在分子图像中,重叠的键必须有不同的方向才能被区分。因此,将0-360个◦均匀地分成60个间隔(60个类别),一次进行60个二元分类任务,以确定在检测位置上是否有一个特定角度的键。在推理阶段,如下图D所示,通过在检测到的键位置搜索大于一定阈值的局部最大值来检测属于特定角度类别的键,对于相反的角度需要非最大抑制,因为对于相反角度的非立体键会有两个正响应,这代表相同的键。
重构算法结合神经网络输出构建graph输出。
通过寻找原子和键热图中的峰点分别来检测原子和键中心。然后,通过寻找沿角度类别轴的局部最大值,可以在每个键中心检测到不同角度类别的键。根据指定的键中心和键角计算出键的长度,对于每个键,可以根据键的中心位置、键角和键长找到两个与这个键相连的原子。原子和键就可以组装成一个分子图。对于其他原子和键的性质,如原子的类型和键的类型,它们可以从原子和键的性质预测图的相应位置推断出来。
二、原子探测
Axy为点(x,y)是关键点的概率。
其中:S为下采样因子,w/h为原始图像。
目标损失函数:
E是Image中原子关键点的数量,α是focal loss的超参数(所有实验α设为2),在推理阶段,热图中的峰是被检测到的原子中心,可通过MaxPooling2D实现。
原子检测后,可以在检测到的原子位置直接预测原子性质,点(x,y)为i类型原子的概率Txy,属性预测损失的计算不仅有ground truth,还有其邻居节点。目标函数为:
三、键探测
预测下采样图像中每个像素成为键中心的概率。将预测目标热图B∈[0,1]中的所有正点为1,一阶邻居设为0.95,其他位置设为0。然后使用惩罚减少的像素级二值焦损失作为训练目标:
F为Image中的键数,τ为focal loss的超参数。
利用热图O∈[0,1]检测B中检测到的候选键位置具有不同角度的键,60个二元分类任务,以确定在指定的角度是否存在键:
将预测目标热图B中的所有像素设为1,将预测目标热图中所有正点的一阶邻域设为0.95,其他位置设为0,然后采用与方程(3)相似的惩罚减少像素focal loss。
在推理阶段,首先通过 2D max-pooling从heatmap B中检测到候选键的位置,将角度分类的局部最大值作为键角。键的性质,如键的类型和键长的确定遵循与原子性质相同的策略。注意,由于键的长度属性是一个回归任务,使用MAE作为损失函数
四、训练过程
1、final loss:在训练过程中自动学习任务权重。
2、Evaluation metrics
其中Ntp、Ndetected和Ngt分别为真阳性、检测和ground truth的数量。
键长任务损失:
M为对应的任务样本总数。
Ntp和Ntn分别为真正数和真负数,N为对应的任务样本的总数。
评估最终的分子结构识别性能,使用了两个指标:准确性(Smiles与预测字符串的)和谷本相似度:
3、训练设置
需要化学结构的二维图像标记像素坐标,以及每个原子和键的相关属性。据我们所知,这种类型的数据集是不可用的,并且手动构建一个包含数千张图像的标记数据集是不可行的。分别用RDKit和Indigo的api分子图像进行描述,可以用于在绘图过程中标记图像。
进一步探索了不同的描述设置:改变原子标记模式,改变键厚,改变分子的取向,改变芳香度标记等。
数据集按8:1:1的比例分成三组(训练、验证和测试)
五、性能分析
每个任务的不同类别下都取得了良好的性能。不同原子的检测性能为>0.98,表明该模型在寻找原子物体方面非常有效。对于普通键,该模型表现良好,但对立体化学键略有下降。在原子电荷方面,ABCNet模型取得了非常令人满意的性能,超过0.99。
与OSRA、MolVec、Img2mol进行对比:
为评估ABCNet模型的泛化能力,进一步利用了UOB数据集作为外部测试数据集。图像的平均分辨率为762×412px,模型可以获得很好的性能,准确率超过95%
训练数据的数量对模型性能的影响。模型被训练来检测分子图像的组成元素(原子和键)。通常,在一幅图像中会有几十个原子和键。因此,训练样本大大超过了图像的数量。
当模型数据量达到80 000时,模型性能达到~95%。随着数据数量的增加,模型性能增长变慢,需要大量广泛的数据样本来获得进一步的微不足道的改进,这是耗时和不必要的。
分子图像噪声的影响。使用二值化图像作为输入,椒盐噪声是模拟实际噪声的最佳方法。
P是背景像素翻转到前景像素的概率,另一个参数为q是相反过程的概率。模型的性能随着噪声水平的提高而下降(上图c),但是如上图d所示,即使在严重的噪声下,模型也能正确识别分子结构的大部分。