tf2 按索引取行和列

TensorFlow 2.x 中的 `tf.tensor` 或者 `tf.data.Dataset` 是张量数据结构，它们并不直接支持像传统二维数组那样通过索引直接获取行和列的操作。然而，你可以使用 `tf.gather`、`tf.squeeze` 和 `tf.expand_dims` 等函数间接实现类似的功能。

例如，如果你有一个形状为 `(batch_size, height, width, channels)` 的四维张量，想要按照行（height方向）取某一行，可以这样做：

# 假设 tensor_data 是你的四维张量
row_index = ...  # 行索引
selected_row = tf.gather(tensor_data, [row_index], axis=1)  # axis=1 表示高度方向

对于取特定列（宽度方向），因为 TensorFlow 通常处理的是通道轴（channels轴），而不是行和列，你需要先转置再取值，然后再转回来：

col_index = ...  # 列索引
if len(tensor_data.shape) == 4:  # 原始是四维张量
    transposed = tf.transpose(tensor_data, perm=[0, 2, 3, 1])  # 转置到 (batch, width, channels, height)
    selected_col = tf.gather(transposed, [col_index], axis=-1)  # axis=-1 表示最后一个轴，即列
    result = tf.transpose(selected_col, perm=[0, 2, 3, 1])  # 再次转回原始顺序
else:
    raise ValueError("This method only works for four-dimensional tensors.")

请注意，这里假设了 `row_index` 和 `col_index` 已经处理过，如果需要从一个更复杂的索引体系中提取元素，你可能需要编写更多的代码来解析索引。