Python tensorflow中的花式索引

我用自定义损失函数实现了一个3D CNN

（Ax'-y）^2

，其中x'是CNN 3D输出的展平和裁剪向量，y是基本真值，a是取x并输出y的线性运算符。因此，我需要一种方法，在计算损失之前，将3D输出展平并使用奇特的索引进行裁剪

以下是我尝试过的： 这是我试图复制的numpy代码

def flatten_crop(img_vol, indices, vol_shape, N):
    """
    :param img_vol: shape (145, 59, 82, N)
    :param indices: shape (396929,)
    """
    nVx, nVy, nVz = vol_shape
    voxels = np.reshape(img_vol, (nVx * nVy * nVz, N), order='F')
    voxels = voxels[indices, :]
    return voxels

我尝试使用

tf.nd\u gather

执行相同的操作，但无法将其推广到任意批量大小。以下是批量大小为1（或单个3D输出）的tensorflow代码：

目前，我的自定义丢失函数中有这段代码，我希望能够推广到任意批量大小。另外，如果您有其他建议来实现这一点（例如自定义层而不是与loss函数集成），我洗耳恭听

谢谢。

请尝试以下代码：

import tensorflow as tf
y_pred = tf.random.uniform((10, 145, 59, 82))
indices = tf.random.uniform((396929,), 0, 145*59*82, dtype=tf.int32)
voxels = tf.reshape(y_pred, (-1, 145 * 59 * 82))   # to flatten and reshape using Fortran-like index order
voxels = tf.gather(voxels, indices, axis=-1)
voxels = tf.transpose(voxels)

请尝试以下代码：

import tensorflow as tf
y_pred = tf.random.uniform((10, 145, 59, 82))
indices = tf.random.uniform((396929,), 0, 145*59*82, dtype=tf.int32)
voxels = tf.reshape(y_pred, (-1, 145 * 59 * 82))   # to flatten and reshape using Fortran-like index order
voxels = tf.gather(voxels, indices, axis=-1)
voxels = tf.transpose(voxels)

---

谢谢！这很有效。我不得不在第四行做了一个改动，以Fortran的顺序进行重塑，在重塑之前添加了一个转置，但除此之外这是完美的。非常感谢！这很有效。我不得不在第4行中做一个修改，以Fortran顺序进行重塑，在重塑之前添加一个转置，但除此之外，这是完美的。