tensorflow内存泄漏或模型只加载不运行

模型只加载不运行"/>

tensorflow内存泄漏或模型只加载不运行

使用tf2模型进行推理的过程中，发现模型的内存占用在逐步增加，甚至会因为OOM被kill掉进程，有时候模型只加载不运行，搜索得到很多五花八门的答案，有些认为是tf2本身的问题，但在使用内存追踪的时候发现，是模型的动态图没有得到释放，而导致这个问题出现的原因，是数据的加载方式存在问题！！！

        mhc_a_batches = list(chunks(mhc_seqs_a, self.batch_size))mhc_b_batches = list(chunks(mhc_seqs_b, self.batch_size))pep_batches = list(chunks(pep_seqs, self.batch_size))assert len(mhc_a_batches) == len(mhc_b_batches)assert len(mhc_a_batches) == len(pep_batches)size = len(mhc_a_batches)# 开始预测preds = []for i in range(size):_preds = self.model([mhc_a_batches[i], mhc_b_batches[i], pep_batches[i]], training = False)preds.extend(_preds.numpy().tolist())return preds

如这段代码，直接使用了list作为模型的输入，尽管tf2也支持numpy的输入格式，但却存在隐患，会产生大量的空tensor！！！

将其改为这样的形式，问题得到解决：

 mhc_seqs_a = tf.convert_to_tensor(mhc_seqs_a, dtype=tf.float32)mhc_seqs_b = tf.convert_to_tensor(mhc_seqs_b, dtype=tf.float32)pep_seqs   = tf.convert_to_tensor(pep_seqs, dtype=tf.float32)assert len(mhc_seqs_a) == len(mhc_seqs_b)assert len(mhc_seqs_a) == len(pep_seqs)ds = tf.data.Dataset.from_tensor_slices((mhc_seqs_a, mhc_seqs_b, pep_seqs)).batch(self.batch_size).prefetch(1)preds = []for x, y, z in ds:_preds = self.model([x,y,z], training=False)preds.extend(_preds.numpy().tolist())return preds

现在可以愉快的进行模型推理了，而且速度比之前要快几倍不止，实测在GPU上提速近30倍，可想而知对于上亿级别的数据，节省的时间多么可观!