О пользовательских операциях в Tensorflow и PyTorch
Я должен реализовать функцию энергии, называемую энергией жесткости, как в уравнении 7 этой статьи здесь.
.
Функция энергии принимает в качестве входных данных две сетки 3D-объектов и возвращает энергию между ними. Первая сетка является исходной сеткой, а вторая сетка является деформированной версией исходной сетки. В грубом псевдокоде вычисления будут выглядеть так:
Перебрать все вершины в исходной сетке.
- Для каждой вершины вычислите ее ковариационную матрицу с соседними вершинами.
- Выполните SVD на вычисленной ковариационной матрице и найдите матрицу вращения вершины.
- Используйте вычисленную матрицу вращения, координаты точки в исходной сетке и соответствующие координаты в деформированной сетке, чтобы вычислить отклонение энергии вершины.
Таким образом, эта функция энергии требует, чтобы я перебирал каждую точку в сетке, а в сетке может быть более 2 тысяч таких точек. В Tensorflow есть два способа сделать это. У меня может быть 2 тензора формы (N,3), один из которых представляет точки источника, а другой - деформированную сетку.
- Делайте это исключительно с помощью тензоров Tensorflow. То есть перебрать элементы приведенных выше тензоров с помощью
tf.gatherи выполнить вычисление в каждой точке, используя только существующие операции TF. Этот метод будет чрезвычайно медленным. Раньше я пытался определить функции потерь, которые повторяют более 1000 точек, и само построение графика занимает слишком много времени, чтобы быть практичным. - Добавьте новый TF OP, как описано в документации по TF здесь. Это включает в себя написание функции в CPP (и Cuda для поддержки GPU) и регистрацию новой OP в TF.
Первый метод прост в написании, но непрактично медленный. Второй способ - больно писать.
Я использую TF в течение 3 лет и никогда раньше не использовал PyTorch, но на данный момент я рассматриваю возможность перехода на него, если он предлагает лучшую альтернативу для таких случаев.
Есть ли в PyTorch способ реализовать такие функции потерь, как легко и работает так же быстро, как и на графическом процессоре. то есть питонический способ написания моих собственных функций потерь, которые работают на графическом процессоре, без какого-либо кода C или Cuda с моей стороны?
Ответы 1
Насколько я понимаю, вы по существу спрашиваете, можно ли эту операцию векторизовать. Ответ — нет, по крайней мере, не полностью, потому что реализация СВД в PyTorch не векторизована.
Если бы вы показали реализацию тензорного потока, это помогло бы понять вашу отправную точку. Я не знаю, что вы имеете в виду, находя матрицу вращения вершины, но я предполагаю, что это можно векторизовать. Это будет означать, что svd является единственной не векторизованной операцией, и вы, возможно, могли бы написать только одну пользовательскую OP, то есть векторизованную svd, что, вероятно, довольно просто, потому что это будет равносильно вызову некоторых библиотечных процедур в цикле в С++.
Я вижу два возможных источника проблем:
- если окрестности
N(i)в уравнении 7 могут иметь существенно разные размеры (это означало бы, что ковариационные матрицы имеют разный размер и векторизация потребовала бы некоторых грязных трюков) - общая проблема работы с сетками и окрестностями может быть сложной. Это врожденное свойство нерегулярных сеток, но в PyTorch есть поддержка разреженных матриц и специальный пакет torch_geometry, который как минимум помогает.