Уравнение теплопроводности в tensorflow
Привет, Хабр! Некоторое время назад увлекся глубоким обучением и стал потихоньку изучать tensorflow. Пока копался в tensorflow вспомнил про свою курсовую по параллельному программированию, которую делал в том году на 4 курсе университета. Задание там формулировалось так:
Линейная начально-краевая задача для двумерного уравнения теплопроводности:
Хотя правильнее было бы назвать это уравнением диффузии.
Задачу тогда требовалось решить методом конечных разностей по неявной схеме, используя MPI для распараллеливания и метод сопряженных градиентов.
Я не специалист в численных методах, пока не специалист в tensorflow, но опыт у меня уже появился. И я загорелся желанием попробовать вычислять урматы на фреймворке для глубокого обучения. Метод сопряженных градиентов реализовывать второй раз уже не интересно, зато интересно посмотреть как с вычислением справится tensorflow и какие сложности при этом возникнут. Этот пост про то, что из этого вышло.
Численный алгоритм
Разностная схема:
Чтобы проще было расписывать, введем операторы:
Явная разностная схема:
В случае явной разностной схемы для вычисления используются значения функции в предыдущий момент времени и не требуется решать уравнение на значения . Однако такая схема менее точная и требует значительно меньший шаг по времени.
Неявная разностная схема:
Перенесем в левую сторону все связанное с , а в правую и домножим на :
По сути мы получили операторное уравнение над сеткой:
что, если записать значения в узлах сетки как обычный вектор, является обычной системой линейных уравнений (). Значения в предыдущий момент времени константы, так как уже рассчитаны.
Для удобства представим оператор как разность двух операторов:
Заменив на нашу оценку , запишем функционал ошибки:
где — ошибка в узлах сетки.
Будем итерационно минимизировать функционал ошибки, используя градиент.
В итоге задача свелась к перемножению тензоров и градиентному спуску, а это именно то, для чего tensorflow и был задуман.
Реализация на tensorflow
Кратко о tensorflow
В tensorflow сначала строится граф вычислений. Ресурсы под граф выделяются внутри tf.Session. Узлы графа — это операции над данными. Ячейками для входных данных в граф служат tf.placeholder. Чтобы выполнить граф, надо у объекта сессии запустить метод run, передав в него интересующую операцию и входные данные для плейсхолдеров. Метод run вернет результат выполнения операции, а также может изменить значения внутри tf.Variable в рамках сессии.
tensorflow сам умеет строить графы операций, реализующие backpropagation градиента, при условии, что в оригинальном графе присутствуют только операции, для которых реализован градиент (пока не у всех).
Сначала код инициализации. Здесь производим все предварительные операции и считаем все, что можно посчитать заранее.
По-хорошему надо было считать значения функции на краях заданными и оптимизировать значения функции только во внутренней области, но с этим возникли проблемы. Способа сделать оптимизируемым только часть тензора не нашлось, и у операции присвоения значения срезу тензора не написан градиент (на момент написания поста). Можно было бы попробовать хитро повозиться на краях или написать свой оптимизатор. Но и просто добавление разности на краях значений функции и краевых условий в функционал ошибки хорошо работает.
Стоит отметить, что метод с адаптивным моментом показал себя наилучшим образом, пусть функционал ошибки и квадратичный.
Вычисление функции: в каждый момент времени делаем несколько оптимизационных итераций, пока не превысим maxiter или ошибка не станет меньше eps, сохраняем и переходим к следующему моменту.
Запуск:
Результаты
Условие как и оригинальное, но без в уравнении:
Что легко правится в коде:
Разницы почти нет, потому что производные имеют большие порядки, чем сама функция.
Условие с одним нагревающимся краем:
Условие с остыванием изначально нагретой области:
Условие с включением нагрева в области:
Рисование гифок
Функция рисования 3D-гифки:
В основной класс добавляем метод, возвращающий U в виде pandas.DataFrame
Функция рисования 2D-гифки:
Стоит отметить, что оригинальное условие без использования GPU считалось 4м 26с, а с использованием GPU 2м 11с. При больших значениях точек разрыв растет. Однако не все операции в полученном графе GPU-совместимы.
- Intel Core i7 6700HQ 2600 МГц,
- NVIDIA GeForce GTX 960M.
Посмотреть, какие операции на чем выполняются, можно с помощью следующего кода:
Это был интересный опыт. Tensorflow неплохо показал себя для этой задачи. Может быть даже такой подход получит какое-то применение — всяко приятнее писать код на питоне, чем на C/C++, а с развитием tensorflow станет еще проще.