Preparing inputs, picking predictions and building masks are all indexing and reshaping tasks.
Basic slicing follows Python and NumPy: x[1], x[:, 1], x[1:, ::2], x[..., 0], and tf.newaxis (alias for None) to insert a dimension.
What NumPy does with fancy indexing, TF does with functions:
tf.gather(x, indices, axis=0)picks slices by index, e.g. rows 2 and 0.tf.gather_nd(x, [[0, 1], [2, 3]])picks individual elements by full coordinates.tf.boolean_mask(x, mask)keeps elements where a boolean mask is true (returns a flat result when the mask covers all dims).tf.where(cond, a, b)chooses element-wise between two tensors, with broadcasting.
Because tensors are immutable, "updating" a slice means building a new tensor, for example with tf.where or tf.tensor_scatter_nd_update.
Reshaping:
tf.reshape(x, [3, -1])regroups elements (one-1is inferred).tf.expand_dims(x, 0)andtf.squeeze(x)add or remove size-1 dimensions.tf.transpose(x, perm=[0, 2, 1])reorders dimensions.tf.concat(list, axis)joins along an existing axis;tf.stack(list, axis)adds a new one.
TensorFlow's image layers use channels-last layout (batch, height, width, channels) by default, unlike PyTorch.
Example
import tensorflow as tf
x = tf.reshape(tf.range(12), [3, 4])
print(x[:, 1].numpy()) # [1 5 9]
print(x[1:, ::2].numpy()) # [[ 4 6] [ 8 10]]
print(tf.gather(x, [2, 0])[:, 0].numpy()) # [8 0] rows 2 and 0
print(tf.boolean_mask(x, x > 8).numpy()) # [ 9 10 11]
print(tf.where(x > 8, x, 0)[2].numpy()) # [ 0 9 10 11]
print(tf.expand_dims(x, 0).shape) # (1, 3, 4)
print(tf.transpose(x).shape) # (4, 3)
print(tf.concat([x, x], axis=0).shape) # (6, 4)
print(tf.stack([x, x]).shape) # (2, 3, 4)
print(tf.squeeze(tf.zeros([1, 3, 1])).shape) # (3,)Remember
- Slicing like NumPy;
tf.gather/boolean_mask/wherefor fancy selection - Updates create new tensors (
tf.where,tensor_scatter_nd_update) - Images are channels-last: (N, H, W, C)
Watch out for
- Expecting
x[[2, 0]]to gather rows as in NumPy; usetf.gather.