Разрезать 4-мерный тензор на 4-мерный тензор меньших субтеноров (вырезать только в двух последних измерениях)

Question

Вопрос аналогичен Разрезать 2d массив на меньшие 2d массивы за исключением того факта, что я использую тензоры (факел) и у меня есть 4D, а не 2D, тензор формы, например. (3, 1, 32, 32) - в моем случае это 3 изображения размером 32х32.

Я хочу разбить каждый тензор формы [i, 0,:,:] на меньшие подмассивы, чтобы выходные данные имели форму, например. (3, 16, 8, 8), где каждый [:, j,:,:] представляет собой небольшой квадрат, вырезанный из исходного изображения. Я не могу найти способ изменить предложенное решение для 4D-тензора.

Я также пытался просто использовать

subx = x.reshape(3, 16, 8, 8)

, но это не меняет его так, как я хочу.

Answer 1

reshape не будет работать для этой цели. Вы можете посмотреть skimage view_as_blocks, где результирующие блоки представляют собой непересекающихся представлений входного массива:

from skimage.util.shape import view_as_blocks
view_as_blocks(a, block_shape=(3,1,8,8)).reshape(3, 16, 8, 8)