网格坐标之间的距离

Question

我有以下代码：

import numpy as np
import tensorflow as tf


with tf.device("cpu:0"):
    sess = tf.InteractiveSession()
    nRows = 10
    nCols = 15
    nNodes = nRows * nCols
    inputDim = 1

    rowIdxs = np.zeros([nRows, nCols])
    colIdxs = np.zeros([nRows, nCols])
    for rowIdx in range(nRows):
        for colIdx in range(nCols):
            rowIdxs[rowIdx, colIdx] = rowIdx
            colIdxs[rowIdx, colIdx] = colIdx

    rowIdxs = np.reshape(rowIdxs, [-1])
    colIdxs = np.reshape(colIdxs, [-1])

    distances = np.zeros([nNodes, nNodes])
    for i in range(nNodes):
        for j in range(nNodes):
            distances[i, j] = ((rowIdxs[i] - rowIdxs[j]) ** 2 + (colIdxs[i] - colIdxs[j]) ** 2)

    print('nRows', nRows, 'nCols', nCols, 'nNodes', nNodes)
    print('distances= ', distances.shape)
    print(distances)

未使用TensorFlow，为方便起见，我已离开导入和会话。

基本上，此代码计算矩阵D(i, j)，其中(i, j)处的元素包含线性化坐标数组中节点i和i之间的欧几里德距离。

即使nRows和nCols的值较小，例如100和100，此代码也很慢。由于simmetry，我只能计算三角矩阵，但问题仍然有效。

这是我多年前写的一个C ++代码的python中的翻译。由于我对TensorFlow和Python的经验都很少，所以我想问一下如何加快计算速度的建议。 TensorFlow是否提供用于计算此缓存矩阵的任何API？我搜索了文档，却找不到任何文档。