Как вычислить матрицу парных расстояний L1 на очень больших изображениях только по соседству?

Я работаю над подходом глубокого обучения для моего проекта. И мне нужно рассчитать матрицу расстояний на 4D Tensor, который будет иметь размер Н х 128 х 64 х 64 (размер партии х каналов х высота х ширина). Матрица расстояний для этого типа тензор будет размером Н х 128 х 4096 х 4096 и будет невозможно установить этот тип тензора в GPU, даже на CPU это потребует много памяти. Итак, я хотел бы рассчитать матрицу расстояний только в некоторых соседних пикселях (скажем, в радиусе 5) и рассмотреть эту прямоугольную матрицу для дальнейшей обработки в нейронной сети. При таком подходе мой матрица расстояний будет размером Н х 128 х 4096 х 61. Это займет меньше памяти по сравнению с матрицей полного расстояния.

Точно, я пытаюсь реализовать Сверточные сети случайных блужданий для семантической сегментации. Эта сеть должна рассчитать парное расстояние L1 для объектов.

Архитектура

Просто добавить этот тип матрицы расстояний обычно рассчитывается для сегментации изображения с помощью спектральной кластеризации.

Например

X = [[a,b],[c,d]]

L1_dist = [ [0, |a-b|, |a-c|, 0],
[|a-b|, 0, 0, |b-d|],
[|a-c|, 0, 0, |c-d| ],
[0,    |b-d|, |c-d|, 0 ]
]

Final_L1_dist = [ [0, |a-b|, |a-c|], // "a" is near to b and c. Including self element i.e. a
[|a-b|, 0, |b-d|], // "b" is near to a and d.
[|a-c|, 0, |c-d| ], // "c" is near to a and d.
[|b-d|, |c-d|, 0 ] // "d" is near to b and c.
]

Буду признателен, если кто-нибудь поможет мне найти эффективный способ вычисления такой матрицы.

Спасибо