tensorflow contrib模块_使用TensorFlow进行高效DNA嵌入

介绍

利用DNA或其他生物序列(如RNA或蛋白质序列)进行深度学习在过去几年里取得了很大进展。大多数从DNA序列中学习的模型使用one-hot编码方案，该方案使用四个通道来表示四个可能的核苷酸A、C、G和T。

长度为N的DNA序列的one-hot编码是(N×4)矩阵，其中列对应于字母A，C，G和T，并且每行恰好具有一个等于一的项，其他条目为零。

任何此类模型中的第一步是以one-hot方案对输入DNA序列进行编码，该方案可以由以下神经网络层处理，无论它们是全连接的，卷积的还是循环的。

通常，one-hot编码被视为在将数据馈送到模型之前发生的预处理步骤。例如，可以对整个训练数据集进行预编码，然后将其存储在磁盘上以用于训练。在其他情况下，模型包含在一些Python预处理层中，该层在Python-land中执行编码，而不是在TensorFlow图中执行。例如，Selene包是一个在PyTorch中处理生物序列的框架，它将编码作为预处理步骤实现，并在Cython中实现。

然而，纯粹从界面设计的角度来看，从DNA序列进行预测的模型应该接受序列作为字符串，而one-hot编码表示应该被视为一个实现细节，不需要对用户可见。这也使得将模型打包为tf.SavedModel并共享或部署更容易，因为该图本身接受DNA序列而无需用户执行他们自己的编码。

编码器只需要执行一个简单的工作：取一个序列，对于每个核苷酸，根据以下映射输出一个向量。

正常情况下，与运行模型本身的成本相比，DNA嵌入将是一项非常廉价的操作，但一个糟糕的嵌入实现可能最终成为一个严重的瓶颈。

事实证明，有很多方法可以使用本机TensorFlow操作来实现DNA one-hot编码，这里我将介绍三种方法。

我将首先介绍我的三个实现，然后我将对它们进行基准测试，看看是否存在相当大的差异。

使用lookup table

使用DNA序列学习在某些方面类似于自然语言处理，通常使用lookup table将字符串键从词汇表映射到整数id。

TensorFlow在tf.contrib.lookup模块中有一个lookup table类。该contrib模块将在即将推出的TensorFlow 2.0中弃用，但可能会有非常接近的替代品。TensorFlow有tf.one_hot函数，它可以将这些整数ID转换为单热嵌入的功能。

以下Python函数将字符串格式的DNA输入映射到整数ID。我现在将省略one-hot 编码步骤，我们可以稍后单独对这两个步骤进行基准测试。

由于我们需要首先将字符串拆分为单个字符，因此我们使用该tf.string_split函数，因此该函数稍微复杂一些。tf.string_split但是，由于返回稀疏张量，我们需要将其转换回密集向量(lookup table只接受密集向量)。

这个函数稍微有点复杂，因为我们需要先将字符串分割成单独的字符，然后使用tfstring_split函数处理这些字符。因为tfstring_split返回一个稀疏张量，所以我们需要将它转换回一个dense向量(lookup table只接受dense向量)。

最后，table.lookup(seq)将结果作为整数id的张量返回。

使用位操作来计算整数索引

使用只有四个键的lookup table似乎有点过火。但是，有什么更简单的方法可以将DNA字母表映射到tf.one_hot函数的索引呢?一种方法是使用基本的位函数直接计算索引。我们需要执行的映射如下:

我们所需要做的就是使用基本的位操作符&、|、^、~、<>找到一个操作序列，将左边的值转换为右边的值。

为了简单起见，我首先在纯Python中实现了以下一个这样的转换：

对于序列中的每个字母，这个Python代码片段首先清除第5位和第7位最低有效位(从右边开始)，然后向右移1位，除了G和T的值交换之外，几乎得到了正确的位模式。因此，剩下的唯一步骤是用2替换3，用3替换2。第三行使用表达式(nt & 1 << 1)作为掩码，只影响右边第二位的值，然后使用xor翻转最右边的位。下表显示了每个输入的逐步转换。

总而言之，此操作将DNA alphabete映射ACGT到索引0,1,2,3。

TensorFlow中此函数的Python实现如下：

def dna_encode_bit_manipulation(seq, name='dna_encode'): with tf.name_scope(name): bytes = tf.decode_raw(seq, tf.uint8) bytes = tf.bitwise.bitwise_and(bytes, ~((1 << 6) | (1 << 4)) bytes = tf.bitwise.right_shift(bytes, 1) mask = tf.bitwise.bitwise_and(bytes, 2) mask = tf.bitwise.right_shift(mask, 1) bytes = tf.bitwise.bitwise_xor(bytes, mask) return bytes

此函数可以替换上面的lookup table，然后可以使用tf.one_hot函数来获取最终编码。由于此

函数仅使用元素运算，因此它也可以在GPU上非常有效地运行。

使用嵌入表

自然语言处理的另一个概念是嵌入表，它将整数id作为输入并输出包含该id的嵌入的向量。在自然语言处理中，这些嵌入是随机初始化和训练的，但我们可以利用相同的工具将我们的序列映射到固定的one-hot编码。

在上面的两种方法中，我们将DNA字母映射到整数0,…,3，以便我们可以使用tf.one_hot函数。但是如果我们直接使用A、C、G和T的整数ASCII码作为索引呢?我们只需要把表做大一些，这样我们就可以使用tf.nn.embedding_lookup函数。

需要分配一个包含84行的嵌入表，因为这是ASCII码的T。

这种方法还有一些其他优点：它可以很容易地适应编码其他字母，如氨基酸序列，或者我们可以使用这种方法来解释IUPAC核苷酸通配符，例如定义R为嘌呤(A或G)或B定义为“除A之外的任何项”(C、G或T)。如果这是我们想要的，我们可以定义以下嵌入表：

基准

为了验证这些方法，我反复使用它们对人类基因组DMD (dystrophine)中最长的基因(约2.24 Mbp)进行one_hot编码。我使用twobitreader模块从hg38 2bit文件(http://hgdownload.soe.ucsc.edu/goldenPath/hg38/bigZips/)中提取序列，计算反补码，因为DMD在负链上:

基准测试使用的软件是Python 3.6.8(Anaconda)，TensorFlow 1.12.0(conda包，支持GPU和MKL扩展)和CUDA 9.2。

我首先运行了dna_encode_lookup_table和dna_encode_bit_manipulation，没有使用tf.one_hot步骤，以便分别对这两个步骤进行基准测试。

我用timeit.repeat(number=10和repeats=10)。对于位操作方法，执行编码10次的10次重复中的最小值为14ms，对于lookup table方法，为2.17s。

dna_encode_embedding_table函数花了38毫秒直接计算最终的one-hot编码。

如果对前两个函数进行基准测试，然后对tf.one_hot应用程序进行基准测试，就会发现总时间主要由从索引中计算one-hot编码决定。位操作方法和one-hot编码的总时间为41 ms，lookup table方法和one-hot编码的总时间为2.18 s。

最后，下图显示了所有结果：

上图为运行时间用于DNA序列的one-hot编码的不同实现。任务是将DMD基因的完整序列(2.24 Mbp)嵌入 10次​​。蓝条仅显示计算整数索引的时间，橙条显示计算整数索引的总时间，然后是one-hot编码。

lookup table显然是一个糟糕的选择，非常低效。在计算整数索引时，位操作非常快，但如果我们包括计算one-hot编码所需的时间，则其大致与嵌入方法一样快。

总的来说，位操作方法具有一些简洁的魅力，但是当需要嵌入IUPAC通配符值时，或者当您使用不同的字母表(如蛋白质序列)时，它不能被使用。在大多数情况下，嵌入表可能是最实用的，因为它可以与任何序列字母表一起使用，并且运行速度与位操作方法一样快。

完整Python示例代码

import numpy as npimport tensorflow as tfimport twobitreaderimport timeitdef tf_dna_encode_lookup_table(seq, name="dna_encode"): """Map DNA string inputs to integer ids using a lookup table."""  with tf.name_scope(name): # Defining the lookup table mapping_strings = tf.constant(["A