在Tensorflow 2中有效管理内存

如何解决在Tensorflow 2中有效管理内存

我需要计算和减少大距离矩阵，其中大> 150,000行。这给我的机器带来了一些时间和内存的复杂性，因此，我试图设置一些代码以在Tensorflow（或PyTorch甚至Numba）中批量运行此代码，以便可以利用GPU的速度。

我遇到的问题是，像下面那样编写一个tf函数的幼稚方法似乎会产生内存泄漏，最终使我的计算机死机。

我已经在github和这里的stackoverflow上探讨了许多相关的TF问题，看来TF似乎有一个长期的基本内存泄漏问题。唯一可行的解​​决方法似乎是多处理，但我无法运行子流程。

我尝试使用下面链接中提供的装饰器，但是tensorflow会生成此错误：https://github.com/tensorflow/tensorflow/issues/37900 或子进程无限期挂起。

c.f。 Clearing Tensorflow GPU memory after model execution

此外，运行子流程会花费更多时间，从而使整个工作几乎毫无意义。因此，我想知道是否有一个更好的方法，我正在丢失。感谢您的任何想法。

def tf_hamming(a,b):
    l = a.shape[1]
    res = tf.math.divide_no_nan(
        tf.reduce_sum(
            tf.cast(
                tf.not_equal(tf.expand_dims(a,axis=1),tf.expand_dims(b,axis=0)),tf.float16),axis=-1),l)
    return res

def diff_matrix_chunk(data,diff_func,chunksize,offset):
    results = np.empty((chunksize,data.shape[0]))

    if offset is None:
        i = 0
    else:
        i = offset

    a = data.values[i:i + chunksize]
    for j in np.arange(0,data.shape[0],chunksize):

        diffs = diff_func(a,b)
        results[:,j:j + chunksize] = diffs.copy()
        del diffs

    res = pd.DataFrame(results,index=data.index[i:i + results.shape[0]],columns=data.index)
    del results
    _ = gc.collect()
    return res

def chunk_reduce(data,reduce_func,chunksize=1000,axis=0,reduce_func_kwargs={}):
    axis = 0
    reduced_data = []
    for i in tqdm(np.arange(0,data.shape[axis],chunksize)):
        diffs = diff_matrix_chunk(data,i)
        reduced_diffs = reduce_func(diffs,**reduce_func_kwargs)
        del diffs
        _ = gc.collect()
        reduced_data.append(reduced_diffs)
    return pd.concat(reduced_data,axis=axis)