如何在TensorFlow中随机采样像素？含动态尺寸场景实现方法

嘿，这个需求我在项目里经常碰到，不管是训练时已知图像尺寸，还是测试时尺寸不确定，都有简洁的实现方式，下面分两种情况给你讲清楚：

一、已知h和w的场景 #

核心思路是先把空间维度（h和w）展平成一个维度，然后在这个维度上随机挑选100个位置。以PyTorch为例：

import torch

# 模拟输入张量，比如batch=2，h=32，w=32，特征维度d=64
input_tensor = torch.randn(2, 32, 32, 64)
b, h, w, d = input_tensor.shape

# 1. 展平空间维度：从(b,h,w,d)转为(b, h*w, d)
flattened = input_tensor.view(b, h * w, d)

# 2. 为每个batch生成100个随机索引（范围0到h*w-1）
# 这里用randint保证每个batch独立采样
random_indices = torch.randint(0, h * w, (b, 100), device=input_tensor.device)

# 3. 用gather提取对应像素，注意要给索引增加一个维度匹配特征维度
sampled_tensor = flattened.gather(
    dim=1,
    index=random_indices.unsqueeze(-1).repeat(1, 1, d)
)

print(sampled_tensor.shape)  # 输出: torch.Size([2, 100, 64])

二、测试阶段未知h和w的场景 #

其实不用纠结h和w的具体数值，直接通过张量的shape动态获取空间维度的总长度就行，同样用PyTorch实现：

import torch

# 测试时输入的h和w不确定，比如这里是batch=3，h=48，w=48，d=128
input_tensor = torch.randn(3, 48, 48, 128)
b, _, _, d = input_tensor.shape  # 忽略h和w，只取batch和特征维度

# 1. 动态展平空间维度：flatten(1,2)会把第1、2维度（h和w）合并
# 等价于input_tensor.view(b, -1, d)，-1会自动计算h*w的数值
flattened = input_tensor.flatten(1, 2)

# 2. 基于展平后的维度长度生成随机索引
random_indices = torch.randint(0, flattened.shape[1], (b, 100), device=input_tensor.device)

# 3. 同样用gather提取目标像素
sampled_tensor = flattened.gather(
    dim=1,
    index=random_indices.unsqueeze(-1).repeat(1, 1, d)
)

print(sampled_tensor.shape)  # 输出: torch.Size([3, 100, 128])

额外注意点 #

• 如果用TensorFlow，思路完全一致：用tf.reshape或tf.keras.layers.Flatten()展平空间维度，tf.random.uniform生成索引，再用tf.gather提取。
• 一定要保证索引和输入张量在同一设备（CPU/GPU）上，不然会出现设备不匹配的报错，所以代码里加上了device=input_tensor.device。
• 如果需要固定采样结果（比如复现实验），可以提前设置随机种子，比如torch.manual_seed(42)。

内容的提问来源于stack exchange，提问作者acloD128