如何对numpy数组指定轴的前n个元素执行roll滚动操作

实现方案 #

你只需提取目标轴对应范围内的切片，执行滚动操作后赋值回原数组即可，numpy的切片为视图操作，整体性能开销极低。

针对你当前只操作第二轴（通道轴，索引为1）第一个元素的需求，代码如下：

import numpy as np
foo = np.random.rand(32,3,16,16)

# 提取第一个通道的所有数据，沿你需要的轴滚动后赋值回去
# 以下为两种常见滚动场景，按需选择即可：
# 场景1：沿图像高度轴（原数组轴2）滚动1步
foo[:, 0, :, :] = np.roll(foo[:, 0, :, :], 1, axis=1)
# 场景2：沿图像宽度轴（原数组轴3）滚动1步
foo[:, 0, :, :] = np.roll(foo[:, 0, :, :], 1, axis=2)

扩展到前n个元素操作 #

如果需要操作第二轴的前n个元素，只需修改切片范围即可，比如操作前2个通道：

# 对前2个通道沿高度轴滚动1步
foo[:, :2, :, :] = np.roll(foo[:, :2, :, :], 1, axis=2)

原代码问题说明 #

你之前使用的np.roll(foo, 1, 1)是对整个数组沿通道轴做整体滚动，会直接移动3个通道的顺序，自然不会只修改第一个通道的内容。

效果验证 #

你可以用小尺寸数组测试效果，确认符合预期：

# 测试用小数组：2张图，2个通道，高2宽2
test = np.arange(2*2*2*2).reshape(2,2,2,2)
print("操作前第一个通道内容：\n", test[:,0,:,:])
# 对第一个通道沿宽轴滚动1步
test[:,0,:,:] = np.roll(test[:,0,:,:], 1, axis=2)
print("操作后第一个通道内容：\n", test[:,0,:,:])
# 验证第二个通道无变化
print("第二个通道内容：\n", test[:,1,:,:])

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