开发: C++知识库 Java知识库 JavaScript Python PHP知识库人工智能区块链大数据移动开发嵌入式开发工具数据结构与算法开发测试游戏开发网络协议系统运维
教程: HTML教程 CSS教程 JavaScript教程 Go语言教程 JQuery教程 VUE教程 VUE3教程 Bootstrap教程 SQL数据库教程 C语言教程 C++教程 Java教程 Python教程 Python3教程 C#教程
数码: 电脑笔记本显卡显示器固态硬盘硬盘耳机手机 iphone vivo oppo 小米华为单反装机图拉丁

-> 人工智能 -> molecular-graph-bert（一） -> 正文阅读

[人工智能]molecular-graph-bert（一）

基于分子图的 BERT 模型，原文：MG-BERT: leveraging unsupervised atomic representation learning for molecular property prediction，原文解析：MG-BERT | 利用无监督原子表示学习预测分子性质 | 在分子图上应用BERT | GNN | 无监督学习(掩蔽原子预训练) | attention，代码：Molecular-graph-BERT，其中缺少的数据以logD7.4例，与上一篇文章处理类似，可以删除 Index 列。代码解析从 pretrain 开始，模型整体框架如下：
请添加图片描述

文章目录

1.pretrain

1.pretrain

os.environ["TF_FORCE_GPU_ALLOW_GROWTH"] = "true"
keras.backend.clear_session()
os.environ['CUDA_VISIBLE_DEVICES'] = "0"
optimizer = tf.keras.optimizers.Adam(1e-4)

small = {'name': 'Small', 'num_layers': 3, 'num_heads': 4, 'd_model': 128, 'path': 'small_weights','addH':True}
medium = {'name': 'Medium', 'num_layers': 6, 'num_heads': 8, 'd_model': 256, 'path': 'medium_weights','addH':True}
medium3 = {'name': 'Medium', 'num_layers': 6, 'num_heads': 4, 'd_model': 256, 'path': 'medium_weights3','addH':True}
large = {'name': 'Large', 'num_layers': 12, 'num_heads': 12, 'd_model': 576, 'path': 'large_weights','addH':True}
medium_balanced = {'name':'Medium','num_layers': 6, 'num_heads': 8, 'd_model': 256,'path':'weights_balanced','addH':True}
medium_without_H = {'name':'Medium','num_layers': 6, 'num_heads': 8, 'd_model': 256,'path':'weights_without_H','addH':False}

arch = medium3      ## small 3 4 128   medium: 6 6  256     large:  12 8 516
num_layers = arch['num_layers']
num_heads =  arch['num_heads']
d_model =  arch['d_model']
addH = arch['addH']


dff = d_model*2
vocab_size =17
dropout_rate = 0.1
model = BertModel(num_layers=num_layers,d_model=d_model,dff=dff,num_heads=num_heads,vocab_size=vocab_size)
train_dataset, test_dataset = Graph_Bert_Dataset(path='data/chem.txt',smiles_field='CAN_SMILES',addH=addH).get_data()

定义优化器和参数，多个参数字典是为了比较不同超参数下模型的效果。根据参数构建 BertModel，载入 Graph_Bert_Dataset 数据

1.1.Graph_Bert_Dataset

"""     
{'O': 5000757, 'C': 34130255, 'N': 5244317, 'F': 641901, 'H': 37237224, 'S': 648962, 
'Cl': 373453, 'P': 26195, 'Br': 76939, 'B': 2895, 'I': 9203, 'Si': 1990, 'Se': 1860, 
'Te': 104, 'As': 202, 'Al': 21, 'Zn': 6, 'Ca': 1, 'Ag': 3}

H C N O F S  Cl P Br B I Si Se
"""

str2num = {'<pad>':0 ,'H': 1, 'C': 2, 'N': 3, 'O': 4, 'F': 5, 'S': 6, 'Cl': 7, 'P': 8, 'Br':  9,
         'B': 10,'I': 11,'Si':12,'Se':13,'<unk>':14,'<mask>':15,'<global>':16}

num2str =  {i:j for j,i in str2num.items()}

class Graph_Bert_Dataset(object):
    def __init__(self,path,smiles_field='Smiles',addH=True):
        if path.endswith('.txt') or path.endswith('.tsv'):
            self.df = pd.read_csv(path,sep='\t')
        else:
            self.df = pd.read_csv(path)
        self.smiles_field = smiles_field
        self.vocab = str2num
        self.devocab = num2str
        self.addH = addH

定义词表 vocab，根据了原子的出现频率决定先后编码序号，可以减少占用空间，vocab 的大小是17，也限定了超参数 vocab_size =17

1.1.1.get_data

def get_data(self):
    data = self.df
    train_idx = []
    idx = data.sample(frac=0.9).index
    train_idx.extend(idx)

    data1 = data[data.index.isin(train_idx)]
    data2 = data[~data.index.isin(train_idx)]

    self.dataset1 = tf.data.Dataset.from_tensor_slices(data1[self.smiles_field].tolist())
    self.dataset1 = self.dataset1.map(self.tf_numerical_smiles).padded_batch(256, padded_shapes=(
        tf.TensorShape([None]),tf.TensorShape([None,None]), tf.TensorShape([None]) ,tf.TensorShape([None]))).prefetch(50)

    self.dataset2 = tf.data.Dataset.from_tensor_slices(data2[self.smiles_field].tolist())
    self.dataset2 = self.dataset2.map(self.tf_numerical_smiles).padded_batch(512, padded_shapes=(
        tf.TensorShape([None]), tf.TensorShape([None, None]), tf.TensorShape([None]),
        tf.TensorShape([None]))).prefetch(50)
    return self.dataset1, self.dataset2

取 90% 的数据作为训练集，tf.data.Dataset.from_tensor_slices 将数据转化为 tensor，map(self.tf_numerical_smiles) 将 smiles 转化为处理后的原子列表，邻接矩阵，未处理原子列表，处理原子标志列表。padded_batch 指定 batch_size，padded_shapes 为 None 会自动取 batch 中长度最长的为准，不足补0（在 vocab 中代表 pad），prefetch 加速运行

1.1.2.tf_numerical_smiles

def tf_numerical_smiles(self, data):
    # x,adjoin_matrix,y,weight = tf.py_function(self.balanced_numerical_smiles,
    #                                           [data], [tf.int64, tf.float32 ,tf.int64,tf.float32])
    x, adjoin_matrix, y, weight = tf.py_function(self.numerical_smiles, [data],
                                                 [tf.int64, tf.float32, tf.int64, tf.float32])

    x.set_shape([None])
    adjoin_matrix.set_shape([None,None])
    y.set_shape([None])
    weight.set_shape([None])
    return x, adjoin_matrix, y, weight

tf.py_function 调用 numerical_smiles，将 smiles 解析为四种数据，set_shape 补全 shape 信息

1.1.3.numerical_smiles

def numerical_smiles(self, smiles):
    smiles = smiles.numpy().decode()
    atoms_list, adjoin_matrix = smiles2adjoin(smiles,explicit_hydrogens=self.addH)
    atoms_list = ['<global>'] + atoms_list
    nums_list =  [str2num.get(i,str2num['<unk>']) for i in atoms_list]
    temp = np.ones((len(nums_list),len(nums_list)))
    temp[1:,1:] = adjoin_matrix
    adjoin_matrix = (1 - temp) * (-1e9)

    choices = np.random.permutation(len(nums_list)-1)[:max(int(len(nums_list)*0.15),1)] + 1
    y = np.array(nums_list).astype('int64')
    weight = np.zeros(len(nums_list))
    for i in choices:
        rand = np.random.rand()
        weight[i] = 1
        if rand < 0.8:
            nums_list[i] = str2num['<mask>']
        elif rand < 0.9:
            nums_list[i] = int(np.random.rand() * 14 + 1)

    x = np.array(nums_list).astype('int64')
    weight = weight.astype('float32')
    return x, adjoin_matrix, y, weight

smiles2adjoin 计算原子列表和邻接矩阵，添加 supernode 后编码为向量，temp 矩阵中，有键相连的是1，(1 - temp) * (-1e9) 后，adjoin_matrix 有键相连的是0，没有键相连的是-1e9
np.random.permutation(len(nums_list)-1) 是 0-len(nums_list)-1的乱序列表，取15%的原子下标，如果15%的原子就取一个原子，+1 应该是为了保证不取到下标0的 supernode，防止被 mask
rand 是概率值，80% 的概率被 mask，10% 的概率被某个原子取代，取代时 rand 的值在(0.8,1)之间，*14+1 后在(11.2,15)，取 int 后相当于被 vocab 中 11-15 所代表的原子取代
最后返回 x 是被 mask 后的训练数据，adjoin_matrix 是原子邻接矩阵，有键相连的是0，没有键相连的是-1e9，y 是weight 是经过处理的原子标志，y 是原始原子列表，即预测目标。示例如下：

import numpy as np
from utils import smiles2adjoin
import tensorflow as tf

str2num = {'<pad>':0 ,'H': 1, 'C': 2, 'N': 3, 'O': 4, 'F': 5, 'S': 6, 'Cl': 7, 'P': 8, 'Br':  9,
         'B': 10,'I': 11,'Si':12,'Se':13,'<unk>':14,'<mask>':15,'<global>':16}

num2str =  {i:j for j,i in str2num.items()}

def numerical_smiles(smiles):
    addH=True
    #smiles = smiles.numpy().decode()
    atoms_list, adjoin_matrix = smiles2adjoin(smiles,explicit_hydrogens=addH)
    atoms_list = ['<global>'] + atoms_list
    nums_list =  [str2num.get(i,str2num['<unk>']) for i in atoms_list]
    temp = np.ones((len(nums_list),len(nums_list)))
    temp[1:,1:] = adjoin_matrix
    adjoin_matrix = (1 - temp) * (-1e9)

    choices = np.random.permutation(len(nums_list)-1)[:max(int(len(nums_list)*0.15),1)] + 1
    y = np.array(nums_list).astype('int64')
    weight = np.zeros(len(nums_list))
    for i in choices:
        rand = np.random.rand()
        weight[i] = 1
        if rand < 0.8:
            nums_list[i] = str2num['<mask>']
        elif rand < 0.9:
            nums_list[i] = int(np.random.rand() * 14 + 1)

    x = np.array(nums_list).astype('int64')
    weight = weight.astype('float32')
    return x, adjoin_matrix, y, weight

smiles='CC(C)OC(=O)C(C)NP(=O)(OCC1C(C(C(O1)N2C=CC(=O)NC2=O)(C)F)O)OC3=CC=CC=C3'

x, adjoin_matrix, y, weight=numerical_smiles(smiles)
x, adjoin_matrix, y, weight
"""
(array([16,  2,  2,  2,  4,  2,  4,  2,  2,  3,  8,  4,  4,  2,  2,  2,  2,
         2,  4,  3,  2,  2,  2,  4,  3,  2,  4,  2, 15,  4,  4,  2,  2,  2,
        15, 15,  2,  1, 15,  1,  1,  1,  1, 15,  1,  1,  1,  1,  1, 15,  1,
         1,  1,  1,  1,  1,  1,  1,  1, 15,  1,  1,  1,  1,  1,  1]),
 array([[-0.e+00, -0.e+00, -0.e+00, ..., -0.e+00, -0.e+00, -0.e+00],
        [-0.e+00, -0.e+00, -0.e+00, ..., -1.e+09, -1.e+09, -1.e+09],
        [-0.e+00, -0.e+00, -0.e+00, ..., -1.e+09, -1.e+09, -1.e+09],
        ...,
        [-0.e+00, -1.e+09, -1.e+09, ..., -0.e+00, -1.e+09, -1.e+09],
        [-0.e+00, -1.e+09, -1.e+09, ..., -1.e+09, -0.e+00, -1.e+09],
        [-0.e+00, -1.e+09, -1.e+09, ..., -1.e+09, -1.e+09, -0.e+00]]),
 array([16,  2,  2,  2,  4,  2,  4,  2,  2,  3,  8,  4,  4,  2,  2,  2,  2,
         2,  4,  3,  2,  2,  2,  4,  3,  2,  4,  2,  5,  4,  4,  2,  2,  2,
         2,  2,  2,  1,  1,  1,  1,  1,  1,  1,  1,  1,  1,  1,  1,  1,  1,
         1,  1,  1,  1,  1,  1,  1,  1,  1,  1,  1,  1,  1,  1,  1]),
 array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
        0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0.,
        1., 1., 0., 0., 1., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 1., 0.,
        0., 0., 0., 0., 0., 0., 1., 0., 1., 0., 0., 0., 0., 0., 0.],
       dtype=float32))
"""

邻接矩阵中 supernode 与所有原子相连，即第一行第一列全为0

1.1.4.smiles2adjoin

def smiles2adjoin(smiles,explicit_hydrogens=True,canonical_atom_order=False):

    mol = Chem.MolFromSmiles(smiles)
    if mol is None:
        print('error')
        #mol = Chem.MolFromSmiles(obsmitosmile(smiles))
        #assert mol is not None, smiles + ' is not valid '

    if explicit_hydrogens:
        mol = Chem.AddHs(mol)
    else:
        mol = Chem.RemoveHs(mol)

    if canonical_atom_order:
        new_order = rdmolfiles.CanonicalRankAtoms(mol)
        mol = rdmolops.RenumberAtoms(mol, new_order)
    num_atoms = mol.GetNumAtoms()
    atoms_list = []
    for i in range(num_atoms):
        atom = mol.GetAtomWithIdx(i)
        atoms_list.append(atom.GetSymbol())

    adjoin_matrix = np.eye(num_atoms)
    # Add edges
    num_bonds = mol.GetNumBonds()
    for i in range(num_bonds):
        bond = mol.GetBondWithIdx(i)
        u = bond.GetBeginAtomIdx()
        v = bond.GetEndAtomIdx()
        adjoin_matrix[u,v] = 1.0
        adjoin_matrix[v,u] = 1.0
    return atoms_list,adjoin_matrix

先验证来自数据库的 smile 是否有效（这里将 obsmitosmile 及 openbabel 的导入删去，对解析影响不大），再计算原子列表和邻接矩阵。atom.GetSymbol() 示例如下：

from rdkit import Chem
from rdkit.Chem.Draw import IPythonConsole
mol=Chem.MolFromSmiles('OC1C2C1CC2')
num_atoms = mol.GetNumAtoms()
for i in range(num_atoms):
    atom = mol.GetAtomWithIdx(i)
    print(atom.GetSymbol(),end='')  #OCCCCC

1.1.5.summary

将 smiles 数据处理为原子列表，邻接矩阵，未处理原子列表，处理原子标志列表。以logD7.4数据为例,，修改 sep 。test_dataset 与 train_dataset 类似，只是 batch_size 不同。get_data 后的情况如下：

"""
class Graph_Bert_Dataset(object):
    def __init__(self,path,smiles_field='Smiles',addH=True):
        if path.endswith('.txt') or path.endswith('.tsv'):
            self.df = pd.read_csv(path,sep='\t') 改为sep=','
        else:
            self.df = pd.read_csv(path)
        self.smiles_field = smiles_field
        self.vocab = str2num
        self.devocab = num2str
        self.addH = addH
"""
from dataset import Graph_Bert_Dataset
addH=True
train_dataset, test_dataset = Graph_Bert_Dataset(path='data/logD.txt',smiles_field='SMILES',addH=addH).get_data()
for (i,(x, adjoin_matrix ,y , char_weight)) in enumerate(train_dataset):
    print("x=\n",x)
    print("adjoin_matrix=\n",adjoin_matrix)
    print("y=\n",y)
    print("char_weight=\n",char_weight)
    if i==2:break
"""
x=
 tf.Tensor(
[[16  5  2 ...  0  0  0]
 [16  6  4 ...  0  0  0]
 [16 15  2 ...  0  0  0]
 ...
 [16  4  2 ...  0  0  0]
 [16  4  2 ...  0  0  0]
 [16 15  2 ...  0  0  0]], shape=(256, 115), dtype=int64)
adjoin_matrix=
 tf.Tensor(
[[[-0. -0. -0. ...  0.  0.  0.]
  [-0. -0. -0. ...  0.  0.  0.]
  [-0. -0. -0. ...  0.  0.  0.]
  ...
  [ 0.  0.  0. ...  0.  0.  0.]
  [ 0.  0.  0. ...  0.  0.  0.]
  [ 0.  0.  0. ...  0.  0.  0.]]

 [[-0. -0. -0. ...  0.  0.  0.]
  [-0. -0. -0. ...  0.  0.  0.]
  [-0. -0. -0. ...  0.  0.  0.]
  ...
  [ 0.  0.  0. ...  0.  0.  0.]
  [ 0.  0.  0. ...  0.  0.  0.]
  [ 0.  0.  0. ...  0.  0.  0.]]

 [[-0. -0. -0. ...  0.  0.  0.]
  [-0. -0. -0. ...  0.  0.  0.]
  [-0. -0. -0. ...  0.  0.  0.]
  ...
  [ 0.  0.  0. ...  0.  0.  0.]
  [ 0.  0.  0. ...  0.  0.  0.]
  [ 0.  0.  0. ...  0.  0.  0.]]

 ...

 [[-0. -0. -0. ...  0.  0.  0.]
  [-0. -0. -0. ...  0.  0.  0.]
  [-0. -0. -0. ...  0.  0.  0.]
  ...
  [ 0.  0.  0. ...  0.  0.  0.]
  [ 0.  0.  0. ...  0.  0.  0.]
  [ 0.  0.  0. ...  0.  0.  0.]]

 [[-0. -0. -0. ...  0.  0.  0.]
  [-0. -0. -0. ...  0.  0.  0.]
  [-0. -0. -0. ...  0.  0.  0.]
  ...
  [ 0.  0.  0. ...  0.  0.  0.]
  [ 0.  0.  0. ...  0.  0.  0.]
  [ 0.  0.  0. ...  0.  0.  0.]]

 [[-0. -0. -0. ...  0.  0.  0.]
  [-0. -0. -0. ...  0.  0.  0.]
  [-0. -0. -0. ...  0.  0.  0.]
  ...
  [ 0.  0.  0. ...  0.  0.  0.]
  [ 0.  0.  0. ...  0.  0.  0.]
  [ 0.  0.  0. ...  0.  0.  0.]]], shape=(256, 115, 115), dtype=float32)
y=
 tf.Tensor(
[[16  5  2 ...  0  0  0]
 [16  6  4 ...  0  0  0]
 [16  4  2 ...  0  0  0]
 ...
 [16  4  2 ...  0  0  0]
 [16  4  2 ...  0  0  0]
 [16  4  2 ...  0  0  0]], shape=(256, 115), dtype=int64)
char_weight=
 tf.Tensor(
[[0. 0. 0. ... 0. 0. 0.]
 [0. 0. 0. ... 0. 0. 0.]
 [0. 1. 0. ... 0. 0. 0.]
 ...
 [0. 0. 0. ... 0. 0. 0.]
 [0. 0. 0. ... 0. 0. 0.]
 [0. 1. 0. ... 0. 0. 0.]], shape=(256, 115), dtype=float32)
x=
 tf.Tensor(
[[16  7  2 ...  0  0  0]
 [16  7  2 ...  0  0  0]
 [16  4  2 ...  0  0  0]
 ...
 [16  4  2 ...  0  0  0]
 [16 15 15 ...  0  0  0]
 [16  6  2 ...  0  0  0]], shape=(256, 132), dtype=int64)
adjoin_matrix=
 tf.Tensor(
[[[-0. -0. -0. ...  0.  0.  0.]
  [-0. -0. -0. ...  0.  0.  0.]
  [-0. -0. -0. ...  0.  0.  0.]
  ...
  [ 0.  0.  0. ...  0.  0.  0.]
  [ 0.  0.  0. ...  0.  0.  0.]
  [ 0.  0.  0. ...  0.  0.  0.]]

 [[-0. -0. -0. ...  0.  0.  0.]
  [-0. -0. -0. ...  0.  0.  0.]
  [-0. -0. -0. ...  0.  0.  0.]
  ...
  [ 0.  0.  0. ...  0.  0.  0.]
  [ 0.  0.  0. ...  0.  0.  0.]
  [ 0.  0.  0. ...  0.  0.  0.]]

 [[-0. -0. -0. ...  0.  0.  0.]
  [-0. -0. -0. ...  0.  0.  0.]
  [-0. -0. -0. ...  0.  0.  0.]
  ...
  [ 0.  0.  0. ...  0.  0.  0.]
  [ 0.  0.  0. ...  0.  0.  0.]
  [ 0.  0.  0. ...  0.  0.  0.]]

 ...

 [[-0. -0. -0. ...  0.  0.  0.]
  [-0. -0. -0. ...  0.  0.  0.]
  [-0. -0. -0. ...  0.  0.  0.]
  ...
  [ 0.  0.  0. ...  0.  0.  0.]
  [ 0.  0.  0. ...  0.  0.  0.]
  [ 0.  0.  0. ...  0.  0.  0.]]

 [[-0. -0. -0. ...  0.  0.  0.]
  [-0. -0. -0. ...  0.  0.  0.]
  [-0. -0. -0. ...  0.  0.  0.]
  ...
  [ 0.  0.  0. ...  0.  0.  0.]
  [ 0.  0.  0. ...  0.  0.  0.]
  [ 0.  0.  0. ...  0.  0.  0.]]

 [[-0. -0. -0. ...  0.  0.  0.]
  [-0. -0. -0. ...  0.  0.  0.]
  [-0. -0. -0. ...  0.  0.  0.]
  ...
  [ 0.  0.  0. ...  0.  0.  0.]
  [ 0.  0.  0. ...  0.  0.  0.]
  [ 0.  0.  0. ...  0.  0.  0.]]], shape=(256, 132, 132), dtype=float32)
y=
 tf.Tensor(
[[16  7  2 ...  0  0  0]
 [16  7  2 ...  0  0  0]
 [16  4  2 ...  0  0  0]
 ...
 [16  4  2 ...  0  0  0]
 [16  4  2 ...  0  0  0]
 [16  6  2 ...  0  0  0]], shape=(256, 132), dtype=int64)
char_weight=
 tf.Tensor(
[[0. 0. 0. ... 0. 0. 0.]
 [0. 1. 0. ... 0. 0. 0.]
 [0. 0. 0. ... 0. 0. 0.]
 ...
 [0. 0. 0. ... 0. 0. 0.]
 [0. 1. 1. ... 0. 0. 0.]
 [0. 0. 0. ... 0. 0. 0.]], shape=(256, 132), dtype=float32)
x=
 tf.Tensor(
[[16  4  2 ...  0  0  0]
 [16  4 15 ...  0  0  0]
 [16  7  2 ...  0  0  0]
 ...
 [16 15  2 ...  0  0  0]
 [16  4  8 ...  0  0  0]
 [16  3  2 ...  0  0  0]], shape=(256, 130), dtype=int64)
adjoin_matrix=
 tf.Tensor(
[[[-0. -0. -0. ...  0.  0.  0.]
  [-0. -0. -0. ...  0.  0.  0.]
  [-0. -0. -0. ...  0.  0.  0.]
  ...
  [ 0.  0.  0. ...  0.  0.  0.]
  [ 0.  0.  0. ...  0.  0.  0.]
  [ 0.  0.  0. ...  0.  0.  0.]]

 [[-0. -0. -0. ...  0.  0.  0.]
  [-0. -0. -0. ...  0.  0.  0.]
  [-0. -0. -0. ...  0.  0.  0.]
  ...
  [ 0.  0.  0. ...  0.  0.  0.]
  [ 0.  0.  0. ...  0.  0.  0.]
  [ 0.  0.  0. ...  0.  0.  0.]]

 [[-0. -0. -0. ...  0.  0.  0.]
  [-0. -0. -0. ...  0.  0.  0.]
  [-0. -0. -0. ...  0.  0.  0.]
  ...
  [ 0.  0.  0. ...  0.  0.  0.]
  [ 0.  0.  0. ...  0.  0.  0.]
  [ 0.  0.  0. ...  0.  0.  0.]]

 ...

 [[-0. -0. -0. ...  0.  0.  0.]
  [-0. -0. -0. ...  0.  0.  0.]
  [-0. -0. -0. ...  0.  0.  0.]
  ...
  [ 0.  0.  0. ...  0.  0.  0.]
  [ 0.  0.  0. ...  0.  0.  0.]
  [ 0.  0.  0. ...  0.  0.  0.]]

 [[-0. -0. -0. ...  0.  0.  0.]
  [-0. -0. -0. ...  0.  0.  0.]
  [-0. -0. -0. ...  0.  0.  0.]
  ...
  [ 0.  0.  0. ...  0.  0.  0.]
  [ 0.  0.  0. ...  0.  0.  0.]
  [ 0.  0.  0. ...  0.  0.  0.]]

 [[-0. -0. -0. ...  0.  0.  0.]
  [-0. -0. -0. ...  0.  0.  0.]
  [-0. -0. -0. ...  0.  0.  0.]
  ...
  [ 0.  0.  0. ...  0.  0.  0.]
  [ 0.  0.  0. ...  0.  0.  0.]
  [ 0.  0.  0. ...  0.  0.  0.]]], shape=(256, 130, 130), dtype=float32)
y=
 tf.Tensor(
[[16  4  2 ...  0  0  0]
 [16  4  2 ...  0  0  0]
 [16  7  2 ...  0  0  0]
 ...
 [16  4  2 ...  0  0  0]
 [16  4  2 ...  0  0  0]
 [16  3  2 ...  0  0  0]], shape=(256, 130), dtype=int64)
char_weight=
 tf.Tensor(
[[0. 0. 0. ... 0. 0. 0.]
 [0. 0. 1. ... 0. 0. 0.]
 [0. 0. 0. ... 0. 0. 0.]
 ...
 [0. 1. 0. ... 0. 0. 0.]
 [0. 0. 1. ... 0. 0. 0.]
 [0. 0. 0. ... 0. 0. 0.]], shape=(256, 130), dtype=float32)
"""

每个 batch 原子数量不一致，取 adjoin_matrix 前5行5列，有键相连为-0.e+00，没有键相连为-1.e+09。pad 的原子之间补的是0，可以和 -0.e+00 区分。char_weight 是被处理原子的标记，每一行是一个 smiles，如果被选中处理(无论是用 mask 代替还是其他操作)为1，否则为

<tf.Tensor: shape=(5, 5), dtype=float32, numpy=
array([[-0.e+00, -0.e+00, -0.e+00, -0.e+00, -0.e+00],
       [-0.e+00, -0.e+00, -0.e+00, -1.e+09, -1.e+09],
       [-0.e+00, -0.e+00, -0.e+00, -0.e+00, -1.e+09],
       [-0.e+00, -1.e+09, -0.e+00, -0.e+00, -0.e+00],
       [-0.e+00, -1.e+09, -1.e+09, -0.e+00, -0.e+00]], dtype=float32)>

1.2.BertModel

class BertModel(tf.keras.Model):
    def __init__(self,num_layers = 6,d_model = 256,dff = 512,num_heads = 8,vocab_size = 17,dropout_rate = 0.1):
        super(BertModel, self).__init__()
        self.encoder = Encoder(num_layers=num_layers,d_model=d_model,num_heads=num_heads,dff=dff,input_vocab_size=vocab_size,maximum_position_encoding=200,rate=dropout_rate)
        self.fc1 = tf.keras.layers.Dense(d_model, activation=gelu)
        self.layernorm = tf.keras.layers.LayerNormalization(-1)
        self.fc2 = tf.keras.layers.Dense(vocab_size)

    def call(self,x,adjoin_matrix,mask,training=False):
        x = self.encoder(x,training=training,mask=mask,adjoin_matrix=adjoin_matrix)
        x = self.fc1(x)
        x = self.layernorm(x)
        x = self.fc2(x)
        return x

经过 encoder 和两个全连接层后输出 shape=(batch_size,input_seq_len，vocab_size) 的 x 用于预测 mask 的原子

1.2.1.Encoder

class Encoder(tf.keras.Model):
    def __init__(self, num_layers, d_model, num_heads, dff, input_vocab_size,
                 maximum_position_encoding, rate=0.1):
        super(Encoder, self).__init__()

        self.d_model = d_model
        self.num_layers = num_layers

        self.embedding = tf.keras.layers.Embedding(input_vocab_size, d_model)
        # self.pos_encoding = positional_encoding(maximum_position_encoding,
        #                                         self.d_model)

        self.enc_layers = [EncoderLayer(d_model, num_heads, dff, rate)
                           for _ in range(num_layers)]

        self.dropout = tf.keras.layers.Dropout(rate)

    def call(self, x, training, mask,adjoin_matrix):
        seq_len = tf.shape(x)[1]
        adjoin_matrix = adjoin_matrix[:,tf.newaxis,:,:]
        # adding embedding and position encoding.
        x = self.embedding(x)  # (batch_size, input_seq_len, d_model)
        x *= tf.math.sqrt(tf.cast(self.d_model, tf.float32))

        x = self.dropout(x, training=training)

        for i in range(self.num_layers):
            x,attention_weights = self.enc_layers[i](x, training, mask,adjoin_matrix)
        return x  # (batch_size, input_seq_len, d_model)

定义 Embedding 层和 num_layers 个 enc_layers 层，将 adjoin_matrix 的 shape 变为 (batch_size,1,seq_len_x,seq_len_x)，是为了之后运算的进行。这里不理解为什么 x 经过 Embedding 后与 $\sqrt{d_model}$ 相乘。

1.2.2.EncoderLayer

class EncoderLayer(tf.keras.layers.Layer):
    def __init__(self, d_model, num_heads, dff, rate=0.1):
        super(EncoderLayer, self).__init__()

        self.mha = MultiHeadAttention(d_model, num_heads)
        self.ffn = point_wise_feed_forward_network(d_model, dff)

        self.layernorm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
        self.layernorm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)

        self.dropout1 = tf.keras.layers.Dropout(rate)
        self.dropout2 = tf.keras.layers.Dropout(rate)

    def call(self, x, training, mask,adjoin_matrix):
        attn_output, attention_weights = self.mha(x, x, x, mask,adjoin_matrix)  # (batch_size, input_seq_len, d_model)
        attn_output = self.dropout1(attn_output, training=training)
        out1 = self.layernorm1(x + attn_output)  # (batch_size, input_seq_len, d_model)

        ffn_output = self.ffn(out1)  # (batch_size, input_seq_len, d_model)
        ffn_output = self.dropout2(ffn_output, training=training)
        out2 = self.layernorm2(out1 + ffn_output)  # (batch_size, input_seq_len, d_model)

        return out2,attention_weights

定义 MultiHeadAttention 和 point_wise_feed_forward_network，最后进入 layernorm2 进行了残差连接

1.2.3.point_wise_feed_forward_network

def point_wise_feed_forward_network(d_model, dff):
  return tf.keras.Sequential([
      tf.keras.layers.Dense(dff, activation=gelu),  # (batch_size, seq_len, dff)tf.keras.layers.LeakyReLU(0.01)
      tf.keras.layers.Dense(d_model)  # (batch_size, seq_len, d_model)
  ])

两个全连接层，先将维度放大为原来的两倍(dff = d_model*2)，在返回原大小

1.2.4.MultiHeadAttention

class MultiHeadAttention(tf.keras.layers.Layer):
    def __init__(self, d_model, num_heads):
        super(MultiHeadAttention, self).__init__()
        self.num_heads = num_heads
        self.d_model = d_model

        assert d_model % self.num_heads == 0

        self.depth = d_model // self.num_heads

        self.wq = tf.keras.layers.Dense(d_model)
        self.wk = tf.keras.layers.Dense(d_model)
        self.wv = tf.keras.layers.Dense(d_model)

        self.dense = tf.keras.layers.Dense(d_model)

    def split_heads(self, x, batch_size):
        """Split the last dimension into (num_heads, depth).
        Transpose the result such that the shape is (batch_size, num_heads, seq_len, depth)
        """
        x = tf.reshape(x, (batch_size, -1, self.num_heads, self.depth))
        return tf.transpose(x, perm=[0, 2, 1, 3])

    def call(self, v, k, q, mask,adjoin_matrix):
        batch_size = tf.shape(q)[0]

        q = self.wq(q)  # (batch_size, seq_len, d_model)
        k = self.wk(k)  # (batch_size, seq_len, d_model)
        v = self.wv(v)  # (batch_size, seq_len, d_model)

        q = self.split_heads(q, batch_size)  # (batch_size, num_heads, seq_len_q, depth)
        k = self.split_heads(k, batch_size)  # (batch_size, num_heads, seq_len_k, depth)
        v = self.split_heads(v, batch_size)  # (batch_size, num_heads, seq_len_v, depth)

        # scaled_attention.shape == (batch_size, num_heads, seq_len_q, depth)
        # attention_weights.shape == (batch_size, num_heads, seq_len_q, seq_len_k)
        scaled_attention, attention_weights = scaled_dot_product_attention(
            q, k, v, mask,adjoin_matrix)

        scaled_attention = tf.transpose(scaled_attention,
                                        perm=[0, 2, 1, 3])  # (batch_size, seq_len_q, num_heads, depth)

        concat_attention = tf.reshape(scaled_attention,
                                      (batch_size, -1, self.d_model))  # (batch_size, seq_len_q, d_model)

        output = self.dense(concat_attention)  # (batch_size, seq_len_q, d_model)

        return output, attention_weights

Transformer 中的多头注意力层，这里的 depth 是原文中的 d_k=d_q=d_v，concat 步骤用并行计算代替，没有使用 3i 个 W 矩阵，而是统一计算，再分成 num_heads 个头计算，其中 scaled_dot_product_attention 略有不同，这里利用了邻接矩阵。

1.2.5.scaled_dot_product_attention

def scaled_dot_product_attention(q, k, v, mask,adjoin_matrix):
    """Calculate the attention weights.
    q, k, v must have matching leading dimensions.
    k, v must have matching penultimate dimension, i.e.: seq_len_k = seq_len_v.
    The mask has different shapes depending on its type(padding or look ahead)
    but it must be broadcastable for addition.

    Args:
      q: query shape == (..., seq_len_q, depth)
      k: key shape == (..., seq_len_k, depth)
      v: value shape == (..., seq_len_v, depth_v)
      mask: Float tensor with shape broadcastable
            to (..., seq_len_q, seq_len_k). Defaults to None.

    Returns:
      output, attention_weights
    """

    matmul_qk = tf.matmul(q, k, transpose_b=True)  # (..., seq_len_q, seq_len_k)

    # scale matmul_qk
    dk = tf.cast(tf.shape(k)[-1], tf.float32)
    scaled_attention_logits = matmul_qk / tf.math.sqrt(dk)

    # add the mask to the scaled tensor.
    if mask is not None:
        scaled_attention_logits += (mask * -1e9)
    if adjoin_matrix is not None:
        scaled_attention_logits += adjoin_matrix

        # softmax is normalized on the last axis (seq_len_k) so that the scores
    # add up to 1.
    attention_weights = tf.nn.softmax(scaled_attention_logits, axis=-1)  # (..., seq_len_q, seq_len_k)

    output = tf.matmul(attention_weights, v)  # (..., seq_len_q, depth_v)

    return output, attention_weights

输入 q，k，v 维度都是 (batch_size, num_heads, seq_len_x, depth)，mask 以如下方式生成：

seq = tf.cast(tf.math.equal(x, 0), tf.float32)
mask = seq[:, tf.newaxis, tf.newaxis, :]

x 的这里 shape 是 (batch_size,seq_len_x)，0 的部分是 pad，adjoin_matrix 的 shape 是 (batch_size,1,seq_len_x,seq_len_x)，mask 的 shape 是 (batch_size,1,1,seq_len_x)，示例如下：

import tensorflow as tf
x=tf.convert_to_tensor([     #batch_size=2,seq_len=3
                        [1,0,3],
                        [0,5,6]
])
seq = tf.cast(tf.math.equal(x, 0), tf.float32)
mask = seq[:, tf.newaxis, tf.newaxis, :]
mask
"""
<tf.Tensor: shape=(2, 1, 1, 3), dtype=float32, numpy=
array([[[[0., 1., 0.]]],


       [[[1., 0., 0.]]]], dtype=float32)>
"""

tf.matmul(q, k, transpose_b=True) 表示的是 (batch_size, num_heads, seq_len_x, depth) 的 q 和 (batch_size, num_heads, depth, seq_len_x) 的 k 相乘，得到 (batch_size, num_heads, seq_len_x, seq_len_x) 的 matmul_qk，scaled_attention_logits 是原文中的 scaled_attention，mask 是 pad 的原子，广播相加后给 pad 的注意力为0。与 adjoin_matrix 广播后相加时，由于有键相连为-0.e+00，没有键相连为-1.e+09，因此没有键相连的原子注意力为0，相当于也被 mask 掉了。具体如下：

请添加图片描述

invisible 是指没有键相连的原子，它们之间的联系较少，对某个原子来说(某列)，与它有键相连的原子应该得到注意力，而没有键相连的原子则不应该得到注意力，应该 mask 掉
将 attention 与 x（原子列表，value）相乘后，原子本身及与它相连的那些原子的值会更大，以此将分子的结构信息编码入了 output，再经过全连接层输出最终编码信息

1.3.run

train_step_signature = [
    tf.TensorSpec(shape=(None, None), dtype=tf.int64),
    tf.TensorSpec(shape=(None, None,None), dtype=tf.float32),
    tf.TensorSpec(shape=(None, None), dtype=tf.int64),
    tf.TensorSpec(shape=(None, None), dtype=tf.float32),
]

train_loss = tf.keras.metrics.Mean(name='train_loss')
train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='train_accuracy')
test_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='test_accuracy')
loss_function = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)

for epoch in range(10):
    start = time.time()
    train_loss.reset_states()

    for (batch, (x, adjoin_matrix ,y , char_weight)) in enumerate(train_dataset):
        train_step(x, adjoin_matrix, y , char_weight)

        if batch % 500 == 0:
            print('Epoch {} Batch {} Loss {:.4f}'.format(
                epoch + 1, batch, train_loss.result()))
            print('Accuracy: {:.4f}'.format(train_accuracy.result()))
            #
            # for x, adjoin_matrix ,y , char_weight in test_dataset:
            #     test_step(x, adjoin_matrix, y , char_weight)
            # print('Test Accuracy: {:.4f}'.format(test_accuracy.result()))
            # test_accuracy.reset_states()
            train_accuracy.reset_states()

    print(arch['path'] + '/bert_weights{}_{}.h5'.format(arch['name'], epoch+1))
    print('Epoch {} Loss {:.4f}'.format(epoch + 1, train_loss.result()))
    print('Time taken for 1 epoch: {} secs\n'.format(time.time() - start))
    print('Accuracy: {:.4f}'.format(train_accuracy.result()))
    model.save_weights(arch['path']+'/bert_weights{}_{}.h5'.format(arch['name'],epoch+1))
    print('Saving checkpoint')

pred 是预测 vocab_size 大小的向量，损失函数为交叉熵，训练后输出训练信息，保存参数

1.3.1.train_step

def train_step(x, adjoin_matrix,y, char_weight):
    seq = tf.cast(tf.math.equal(x, 0), tf.float32)
    mask = seq[:, tf.newaxis, tf.newaxis, :]
    with tf.GradientTape() as tape:
        predictions = model(x,adjoin_matrix=adjoin_matrix,mask=mask,training=True)
        loss = loss_function(y,predictions,sample_weight=char_weight)
        gradients = tape.gradient(loss, model.trainable_variables)
        optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    train_loss.update_state(loss)
    train_accuracy.update_state(y,predictions,sample_weight=char_weight)