PINN или нейросети, знающие физику
Что такое PINN и какова область их применения?
PINN появились сравнительно недавно (в 2019 году была опубликована работа Raissi), но уже активно применяются для некоторых задач физики. Отличительная особенность данных нейросетей состоит в том, что в Loss-функцию включены невязки от уравнений, которые описывают рассматриваемый физический процесс. Вход такой нейросети — это координаты (пространственные или временные, в зависимости от задачи). И еще одна особенность — для обучения не требуется таргетов, так как, повторюсь, в Loss минимизируется отклонение предсказанных значений от уравнений.
Можно сказать, что PINN — это замена численному моделированию и тогда, возникает вопрос: «А нужны ли нейросети там, где хорошо работают численные методы». Но не все так просто. Представьте, что Вы или Ваш коллега провел эксперимент, например, измерил скорости частиц в потоке жидкости, или получил точечные измерения температуры. Если Вы действительно имеете опыт проведения эксперимента, наверняка знаете, что экспериментальные данные далеко не идеальны и могут принести много головной боли при обработке. А теперь представим, что все же Вы эту обработку провели, получили датасет из эксперимента, и теперь хотите эти данные использовать в уравнениях, чтобы получить другие параметры течения. Например, измерили скорость, а из уравнений гидродинамики, хотите получить давление. Или другими словами, провести ассимиляцию данных, говоря на современном околонаучным языком. Численное моделирование в таком случае можно дать сбой, потому что даже тщательно отфильтрованные данные могут быть шумными (особенно если от них требуется брать производные, а если еще и вторые, то совсем, шляпа). Или их может быть мало (например, температуру измеряли термопарой в нескольких точках). В данном случае, вроде, эксперимент есть и потенциально восстанавливать одни величины по другим можно, решая уравнения. И тут на помощь могут прийти PINN. Потому что они работают иначе, чем численные методы. Они не используют схем переноса, а параметры нейросети минимизируются в выбранных точках.
Кроме того, в PINN не используются традиционные численные производные, а используется такой инструмент — автоматическое дифференцирование. Ведь, напомню, PINN существуют, чтобы минимизировать невязки от дифференциальных уравнений. Постараюсь объяснить «наивно», что такое «автоматические» производные. PINN — полносвязная нейросеть, то есть несколько слоев с определенным количеством нейронов на каждом слое. При обратном проходе по входным параметрам (напомню, что здесь входные данные — это пространственные координаты, время) берутся производные. А теперь задумаемся: ведь это не что иное, как обычные производные, только, можно сказать, «производная» в точке, в пределе. Обычная численная производная — это разности, а здесь — это просто вложенные в нейросеть весовые коэффициенты. Мы же можем использовать эти производные, как необходимые нам. И вторые производные, и третьи… Да, можем. И вот мощное преимущество PINN. Вспомним, что экспериментальные данные — это часто шумные, разреженные данные. А нейросеть, при правильном обучении, может выучить эти данные и предоставит нам автоматические производные в нужным нам точках.
В этом состоит идея использования нейросетей для ассимиляции данных.
Простой пример с гармоническим осциллятором
В данном обзоре я хочу на простом примере показать, как обучить нейросеть решать уравнение для обычного гармонического осциллятора (грузика массы m на пружине жесткостью k) и продемонстрировать, что такое «автоматические» производные, указать на некоторые особенности PINN. Возьмем самый обычный грузик на пружине. Такой пример изучается еще в школе, и хорошо известен. На рис. 1 изображена знакомая схема.
Рис. 1. Схема грузика на пружинке, совершающего гармонические колебания
Уравнение, которым будем физически «информировать» нашу PINN, то есть уравнение движения гармонического осциллятора:
Пускай грузик будет выведен из начального положение равновесия () в положение
и отпущен без начальной скорости:
. Как следует из уравнения движения, .
. Решая данное уравнение, получаем закон движения грузика:
Это и хотелось бы получить после обучения нейросети.
Архитектура нейросети
Теперь самое интересное — как создать PINN и обучить ее? На самом деле, ничего сложного здесь нет. Очень важны такие гиперпараметры, как количество слоев и количество нейронов на каждом слое (то есть глубина нейросети), так как это влияет на время обучения и естественно, на результат. Однозначных ответов, почему надо выбрать конкретное количество нейронов и слоев в разных статьях я не нашла, поэтому тут есть некоторый элемент подбора. Хочу отметить, что лучше начинать с маленького количества слоев (скажем, 5) и маленького количества нейронов (например, по 20 на каждом слое). Конечно, это зависит от количества выходов и входов нейросети. В данном примере я буду показывать результат с нейросетью из 4 скрытых слоев и 20 нейронов на каждом слое. Между слоями используется функция активации гиперболический тангенс, на выходе последнего слоя также есть гиперболический тангенс, чтобы выходы нейросети находились в пределах от -1 до 1. Предварительно я выбрала , чтобы удовлетворять этому условию, но всегда можно сделать надлежащую нормировку, поэтому здесь не так важно, какое начальное условие выбрать.
выбрано для простоты.
Рис. 2. Схема полносвязной нейросети (PINN)
Далее я приведу части кода. Использовалась библиотека pytorch.
В качестве оптимизатора был выбран L-BFGS (о методе можно прочесть здесь). Чаще всего в статьях встречается либо этот оптимизатор, либо использование двух оптимизаторов: предварительное обучение с помощью Adam, затем L-BFGS. Так как задача достаточно простая, я решила не перегружать обучение и использовала только последний оптимизатор.
Вот параметры задачи: возьмем грузик с собственной частотой колебаний (
) и рассмотрим промежуток
, то есть наш грузик должен будет совершить два колебания. Тут я намеренно опускаю все размерности, будет считать, что все величины предварительно обезразмерены.
Собственно, вот сама PINN:
class simpleModel(nn.Module):
def __init__(self,
hidden_size=20):
super().__init__()
self.layers_stack = nn.Sequential(
nn.Linear(1, hidden_size),
nn.Tanh(),
nn.Linear(hidden_size, hidden_size), #1
nn.Tanh(),
nn.Linear(hidden_size, hidden_size), #2
nn.Tanh(),
nn.Linear(hidden_size, hidden_size), #3
nn.Tanh(),
nn.Linear(hidden_size, hidden_size), #4
nn.Tanh(),
nn.Linear(hidden_size, 1),
nn.Tanh(),
)
def forward(self, x):
return self.layers_stack(x)
А вот сам процесс обучения:
steps=1000
pbar = tqdm(range(steps), desc='Training Progress')
t = (torch.linspace(0, 1, 100).unsqueeze(1)).to(device)
t.requires_grad = True
metric_data = nn.MSELoss()
writer = SummaryWriter()
optimizer = torch.optim.LBFGS(model.parameters(), lr=0.1)
def train():
for step in pbar:
def closure():
optimizer.zero_grad()
loss = pdeBC(t)
loss.backward()
return loss
optimizer.step(closure)
if step % 2 == 0:
current_loss = closure().item()
pbar.set_description("Step: %d | Loss: %.6f" %
(step, current_loss))
writer.add_scalar('Loss/train', current_loss, step)
train()
writer.close()
Пожалуй, самое главное здесь, это функция pdeBC ():
def pdeBC(t):
out = model(t).to(device)
f1 = pde(out, t)
inlet_mask = (t[:, 0] == 0)
t0 = t[inlet_mask]
x0 = model(t0).to(device)
dx0dt = torch.autograd.grad(x0, t0, torch.ones_like(t0), create_graph=True, \
retain_graph=True)[0]
loss_bc = metric_data(x0, x0_true) + \
metric_data(dx0dt, dx0dt_true.to(device))
loss_pde = metric_data(f1, torch.zeros_like(f1))
loss = 1e3*loss_bc + loss_pde
metric_x = metric_data(out, x0_true * torch.sin(omega*t + torch.pi / 2))
metric_x0 = metric_data(x0, x0_true)
metric_dx0dt = metric_data(dx0dt, dx0dt_true.to(device))
acc_metrics = {'metric_x': metric_x,
'metric_x0': metric_x0,
'metric_dx0dt': metric_dx0dt,
}
metrics = {'loss': loss,
'loss_bc': loss_bc,
'loss_pde': loss_pde,
}
for k, v in metrics.items():
run[k].append(v)
for k, v in acc_metrics.items():
run[k].append(v)
return loss
Вместе с самим уравнением:
def pde(out, t, nu=2):
omega = 2 * torch.pi * nu
dxdt = torch.autograd.grad(out, t, torch.ones_like(t), create_graph=True, \
retain_graph=True)[0]
d2xdt2 = torch.autograd.grad(dxdt, t, torch.ones_like(t), create_graph=True, \
retain_graph=True)[0]
f = d2xdt2 + (omega ** 2) * out
return f
Обращаю внимание, что torch.autograd.grad — это и есть автоматическая производная. И, как можно заметить, . То есть Loss-функция состоит из двух частей: невязки от уравнений и отклонений от граничных условий. Более сложные PINN включают в себя еще больше частей, например, как отмечалось в введении, если есть экспериментальные данные, можно включить часть, связанную с отклонениями от этих экспериментальных данных. И важно правильно подобрать веса, чтобы нейросеть не стремилась минимизировать только какую-то одну часть. В своем примере я подобрала вес 1000 перед перед
.
Привожу график обучения:
Рис. 3. График обучения
А вот результат предсказанных нейросетью значений и аналитических для самой координаты и полученной с помощью torch.autograd.grad скорости
:
Рис. 4. Результат для координаты и скорости, полученной с помощью автоматического дифференцирования.
Получилось достаточно неплохо (а именно, L2-норма отклонения предсказанного результата от аналитического решения составляет менее 2%).
Заключение. Мы разобрали простой пример решения задачи гармонического осциллятора с помощью PINN. Хотелось показать функционал PINN и инструмент автоматического дифференцирования. Существуют намного более сложные примеры с включением экспериментальных данных, где геометрия трехмерная и задачи нестационарные, но об этом в следующий раз :)
P.S. Код с сохраненными весами модели доступен в репозитории.