如何高效实现PIL图像随机索引位置的数值累加？

原代码的核心问题是Python双重循环的开销过大，且笛卡尔积式的索引生成会导致大量重复计算，针对大尺寸图像（张量）可以用以下两种优化方案：

方案1：无重复随机像素批量加值（推荐） #

直接生成指定数量的唯一像素坐标，通过PyTorch向量化操作一次性完成加值，避免重复计算：

import torch

# 假设cat是shape为(H, W, C)的PyTorch张量，cost为要添加的数值（标量或匹配通道数的张量）
H, W, C = cat.shape

# 计算要处理的像素数量（与原逻辑的笛卡尔积数量一致：25%行 × 25%列）
total_pixels = int(H * 0.25) * int(W * 0.25)

# 生成所有像素的线性索引，随机选择指定数量
all_indices = torch.arange(H * W)
selected_idx = torch.randperm(H * W)[:total_pixels]

# 转换为二维坐标
idx_r = selected_idx // W
idx_c = selected_idx % W

# 批量加值
cat[idx_r, idx_c, :] += cost

方案2：保留原逻辑的有放回采样（加速版） #

如果需要和原代码一样允许同一位置被多次加值（有放回采样），可以用meshgrid生成所有坐标对后批量操作：

import torch
import random

H, W = cat.shape[:2]

# 生成有放回的行、列索引（与原逻辑一致）
idx_r = torch.tensor(random.choices(range(H), k=int(H*0.25)))
idx_c = torch.tensor(random.choices(range(W), k=int(W*0.25)))

# 生成所有笛卡尔积坐标对
r_grid, c_grid = torch.meshgrid(idx_r, idx_c, indexing='ij')

# 展平后批量加值
cat[r_grid.flatten(), c_grid.flatten(), :] += cost

为什么效率更高？ #

• 完全移除了Python层面的循环，改用PyTorch底层优化的张量操作（CPU下是C++实现，GPU下是CUDA加速），效率提升几个数量级。
• 避免了重复索引的冗余计算（方案1），即使保留有放回逻辑，也是通过向量化方式完成，远快于Python循环。

注意事项 #

• 如果处理的是PIL图像，需先转换为PyTorch张量：cat = torch.tensor(np.array(pil_image))，处理完成后再转回PIL：Image.fromarray(cat.numpy().astype(np.uint8))。
• 确保cost的形状与通道数匹配，比如RGB图像的cost可以是标量或形状为(3,)的张量。

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