Несколько месяцев назад у меня было похожее требование, в котором, учитывая 2d точку, я искал из массива отрезков линии, который был ближе всего к точке. Я изо всех сил пытался принудить tenorflowjs выполнить это эффективно, и в конце концов наткнулся на gpu.js, который больше ориентирован на компиляцию пользовательских функций ядра GPU.
В приведенном ниже примере, который я подготовил, у меня есть массивы, представляющие 11 (X, Y) координат, и еще одна пара массивов, представляющих 5 (X, Y) координат. Результатом будет матрица 11x5, которая вычисляет каждое расстояние между обоими наборами точек. Ключевой функцией является «ядро», которое компилируется gpu.js для работы на ядре графического процессора и, по существу, вычисляет расстояние между парой точек, полученных как из 11 координат, так и из 5 координат. Теоретически эта функция ядра будет распределена по многим ядрам графического процессора для повышения производительности. Т.е. в этом случае выполнить все 55 сразу. (Я говорю «в теории», потому что gpu.js, как я понимаю, использует функцию карты шейдеров webGL, и я не совсем уверен в том, какое влияние слои стека виртуализации задействуют в стеке, что приводит к тому, что ядра GPU фактически выполняют эту работу. ..)
Результатом является матрица 11x5, содержащая расстояние от каждой комбинации пар точек, после чего эта матрица 11x5 затем передается по конвейеру в «kernelMin», который будет немного медленнее, поскольку он просматривает результаты в поисках минимального значения, а также захват индекса минимального значения. При этом должно быть 11 параллельных ядер GPU, работающих над попыткой найти, какая из 5 координат является ближайшей.
const kernel = gpu.createKernel(function(x0, y0, x1, y1) {
let dx = x1[this.thread.y][0] - x0[0][this.thread.x];
let dy = y1[this.thread.y][0] - y0[0][this.thread.x];
return Math.sqrt(dx * dx + dy * dy);
}).setPipeline(true).setOutput([11,5]);
const result1 = kernel(
GPU.input(
new Float32Array([0,10,20,30,40,50,60,70,80,90,100]),
[11,1]
),
GPU.input(
new Float32Array([100,100,100,100,100,100,100,100,100,100,100]),
[11,1]
),
GPU.input(
new Float32Array([0,30,50,70,100]),
[1,5]
),
GPU.input(
new Float32Array([10,10,10,10,10]),
[1,5]
)
);
console.log(result1.toArray());
const kernelMin = gpu.createKernel(function(a) {
let minVal = 1000000;
let minIdx = 0;
for (let y = 0; y < 5; y++) {
if (a[y][this.thread.x] < minVal) {
minVal = a[y][this.thread.x];
minIdx = y;
}
}
return [minVal,minIdx];
}).setOutput([11]);
const result2 = kernelMin(result1);
console.log(result2);
Окончательный вывод ...
0: Float32Array(2) [90, 0]
1: Float32Array(2) [90.55384826660156, 0]
2: Float32Array(2) [90.55384826660156, 1]
3: Float32Array(2) [90, 1]
4: Float32Array(2) [90.55384826660156, 1]
5: Float32Array(2) [90, 2]
6: Float32Array(2) [90.55384826660156, 2]
7: Float32Array(2) [90, 3]
8: Float32Array(2) [90.55384826660156, 3]
9: Float32Array(2) [90.55384826660156, 4]
10: Float32Array(2) [90, 4]
Обратите внимание, что для ясности я жестко закодировал размеры матрицы в примере. Gpu.js явно принимает переменные. Кроме того, в вашем случае, в зависимости от размера ваших матриц, вам может потребоваться разбить проблему на куски в зависимости от объема ОЗУ графического процессора, необходимого для размещения полной кросс-матрицы расстояний ...
Я понимаю, что это не tenorflowjs, но надеюсь, что это поможет.
РЕДАКТИРОВАТЬ - через TensorFlow.JS
Потратил несколько минут на портирование на tenorflow.js. Основная идея заключается в построении матриц значений x и y при подготовке к выполнению расчетов по массе.
const x0 = tf.tensor1d([0,10,20,30,40,50,60,70,80,90,100]);
const y0 = tf.tensor1d([100,100,100,100,100,100,100,100,100,100,100]);
const x1 = tf.tensor1d([0,30,50,70,100]);
const y1 = tf.tensor1d([10,10,10,10,10]);
const x0mat = x0.tile([5]).reshape([5,11]);
const y0mat = y0.tile([5]).reshape([5,11]);
const x1mat = x1.tile([11]).reshape([11,5]).transpose();
const y1mat = y1.tile([11]).reshape([11,5]).transpose();
x0mat.print();
x1mat.print();
const xDeltas = x1mat.squaredDifference(x0mat);
y0mat.print();
y1mat.print();
const yDeltas = y1mat.squaredDifference(y0mat);
const distance = xDeltas.add(yDeltas).sqrt();
distance.print();
distance.argMin(1).print();
distance.min(1).print();
С результатами ...
Tensor - x0mat
[[0, 10, 20, 30, 40, 50, 60, 70, 80, 90, 100],
[0, 10, 20, 30, 40, 50, 60, 70, 80, 90, 100],
[0, 10, 20, 30, 40, 50, 60, 70, 80, 90, 100],
[0, 10, 20, 30, 40, 50, 60, 70, 80, 90, 100],
[0, 10, 20, 30, 40, 50, 60, 70, 80, 90, 100]]
Tensor - x1mat
[[0 , 0 , 0 , 0 , 0 , 0 , 0 , 0 , 0 , 0 , 0 ],
[30 , 30 , 30 , 30 , 30 , 30 , 30 , 30 , 30 , 30 , 30 ],
[50 , 50 , 50 , 50 , 50 , 50 , 50 , 50 , 50 , 50 , 50 ],
[70 , 70 , 70 , 70 , 70 , 70 , 70 , 70 , 70 , 70 , 70 ],
[100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100]]
Tensor - y0mat
[[100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100],
[100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100],
[100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100],
[100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100],
[100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100]]
Tensor - y1mat
[[10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10],
[10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10],
[10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10],
[10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10],
[10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10]]
Tensor - distance
[[90 , 90.5538483 , 92.1954422 , 94.8683319, 98.4885788 , 102.9562988, 108.1665344, 114.01754 , 120.415947 , 127.2792206, 134.5362396],
[94.8683319 , 92.1954422 , 90.5538483 , 90 , 90.5538483 , 92.1954422 , 94.8683319 , 98.4885788, 102.9562988, 108.1665344, 114.01754 ],
[102.9562988, 98.4885788 , 94.8683319 , 92.1954422, 90.5538483 , 90 , 90.5538483 , 92.1954422, 94.8683319 , 98.4885788 , 102.9562988],
[114.01754 , 108.1665344, 102.9562988, 98.4885788, 94.8683319 , 92.1954422 , 90.5538483 , 90 , 90.5538483 , 92.1954422 , 94.8683319 ],
[134.5362396, 127.2792206, 120.415947 , 114.01754 , 108.1665344, 102.9562988, 98.4885788 , 94.8683319, 92.1954422 , 90.5538483 , 90 ]]
Tensor - argMin of distance
[0, 3, 5, 7, 10]
Tensor - min of distance
[90, 90, 90, 90, 90]
Код разбит на этапы, чтобы показать основную концепцию. Я уверен, что это может быть сжато и оптимизировано далее.