Обучите нейронную сеть ОДУ

Открыть Live Script

Этот пример показывает, как обучить нейронную сеть с нейронными обыкновенными дифференциальными уравнениями (ОДУ), чтобы узнать динамику физической системы.

Нейронные ОДУ [1] являются операциями глубокого обучения, заданными решением обыкновенного дифференциального уравнения (ОДУ). Более конкретно, нейронная ОДУ является слоем, который может использоваться в любой архитектуре и, учитывая вход, определяет его выход как численное решение следующей ОДУ:

$y^{'} = f (t, y, θ)$

для временного горизонта $(t_{0}, t_{1})$ и с начальным условием $y (t_{0}) = y_{0}$ . Размер правой руки $f (t, y, θ)$ ОДУ зависит от набора обучаемых параметров $θ$ , которые изучаются в процессе обучения. В этом примере $f (t, y, θ)$ моделируется с помощью dlnetwork объект, который встраивается в пользовательский слой. Как правило, начальное условие $y_{0}$ является либо входом всей сети, как в случае с этим примером, либо является выходом предыдущего слоя.

Этот пример показывает, как обучить нейронную сеть с нейронными обыкновенными дифференциальными уравнениями (ОДУ), чтобы узнать динамику $x$ данной физической системы, описанной в следующей ОДУ:

$x^{'} = A x$ ,

где $A$ является матрицей 2 x 2.

В этом примере ОДУ, которая определяет модель, решается численно с помощью явного метода Рунге-Кутты (RK) 4-го порядка [2] в переднем пассе. Обратный проход использует автоматическую дифференциацию, чтобы узнать обучаемые параметры $θ$ , путем обратного распространения через каждую операцию решателя RK.

Пример затем показывает, как использовать обученную функцию $f (t, y, θ)$ как правая сторона модели в ode45 для вычисления решения той же модели из дополнительных начальных условий.

Синтезируйте данные целевой динамики

Задайте целевую динамику как линейную модель ОДУ.

x0 = [2; 0];
A = [-0.1 -1; 1 -0.1];

Найдите численное решение задачи выше. Создайте набор точек данных для использования для обучения с помощью функции ode45.

numTimeSteps = 4000;
T = 15;
odeOptions = odeset('RelTol', 1.e-7, 'AbsTol', 1.e-9);
t = linspace(0, T, numTimeSteps);
[~, x] = ode45(@(t,y) A*y, t, x0, odeOptions);
x = x';

Определение сети

Задайте dlnetwork объект, который используется в качестве Learnable свойство в пользовательском слое neuralOdeLayer.

Для каждого наблюдения neuralOdeInternalDlnetwork принимает вектор длины inputSize, размер решения ОДУ, как вход, он увеличивает его так, чтобы иметь длину hiddenSize и затем сжимает его снова, чтобы иметь длину outputSize. Объект neuralOdeInternalDlnetwork определяет правую сторону $f (t, y, θ)$ ОДУ, которая будет решена, и обучаемые параметры $θ$ являются ли знания neuralOdeInternalDlnetwork. Вы можете сделать архитектуру нейронных одов более выразительной, например, увеличив скрытый размер или количество слоев, которые определяют neuralOdeInternalDlnetwork.

hiddenSize = 30;
inputSize = size(x,1);
outputSize = inputSize;

neuralOdeLayers = [
    fullyConnectedLayer(hiddenSize)
    tanhLayer
    fullyConnectedLayer(hiddenSize)
    tanhLayer
    fullyConnectedLayer(outputSize)
    ];

neuralOdeInternalDlnetwork = dlnetwork(neuralOdeLayers,'Initialize',false);
neuralOdeInternalDlnetwork.Learnables

ans=6×3 table
    Layer     Parameter       Value    
    ______    _________    ____________

    "fc_1"    "Weights"    {0×0 double}
    "fc_1"    "Bias"       {0×0 double}
    "fc_2"    "Weights"    {0×0 double}
    "fc_2"    "Bias"       {0×0 double}
    "fc_3"    "Weights"    {0×0 double}
    "fc_3"    "Bias"       {0×0 double}

Задайте слой нейронной ОДУ, который имеет dlnetwork объект как настраиваемый параметр. Пользовательский слой neuralODELayer принимает мини-пакет данных, который представляет начальные условия и выводит предсказанную динамику численным решением задачи $y^{'} = f (t, y, θ)$ , где $f (t, y, θ)$ моделируется neuralOdeInternalDlnetwork.

The value от neuralOdeInternalTimesteps определяет количество временных интервалов решателя Runge-Kutta ОДУ, внутренне используемых в пользовательском слое neuralOdeLayer. Значение dt - размер шага, используемый во внутреннем методе RK.

neuralOdeInternalTimesteps = 40;
dt = t(2);
neuralOdeLayerName = 'neuralOde';

customNeuralOdeLayer = neuralOdeLayer(neuralOdeInternalDlnetwork,neuralOdeInternalTimesteps,dt,neuralOdeLayerName);

Объявление внешнего dlnetwork объект и инициализация сети. Это также инициализирует объект dlnetwork в пользовательском neuralOdeLayer.

dlnet=dlnetwork(customNeuralOdeLayer,'Initialize',false);
dlnet = initialize(dlnet, dlarray(ones(inputSize,1),'CB'));

Настройка опций обучения

Задайте опции для оптимизации ADAM.

gradDecay = 0.9;
sqGradDecay = 0.999;
learnRate = 0.001;

Обучите для 400 итераций с мини-пакетным размером 50.

numIter = 1500;
miniBatchSize = 200;

Инициализируйте график процесса обучения.

plots = "training-progress";
lossHistory = [];

Каждые 50 итераций решают с помощью ode45 полученную динамику, выраженную как внутренний объект dlnetwork в объекте neuralOdeLayer и отображают его с основной истиной в фазовой схеме, это показывает путь обучения обученной динамики.

plotFrequency = 50;

Обучите сеть с помощью пользовательского цикла обучения

Инициализируйте averageGrad и averageSqGrad параметры для решателя ADAM.

averageGrad = [];
averageSqGrad= [];

Для каждой итерации:

создать мини-пакет данных из синтезированных данных.
оцените градиенты модели и потери с помощью dlfeval и функции modelGradients.
обновляйте параметры сети с помощью adamupdate функция.

if plots == "training-progress"
    figure(1)
    clf
    title('Training Loss');
    lossline = animatedline;
    xlabel('Iteration')
    ylabel("Loss")
    grid on
end
numTrainingTimesteps = numTimeSteps;
trainingTimesteps = 1:numTrainingTimesteps;

start = tic;

for iter=1:numIter
    
    % Create batch 
    [dlx0, targets] = createMiniBatch(numTrainingTimesteps, neuralOdeInternalTimesteps, miniBatchSize, x);
    
    % Evaluate network and compute gradients 
    [grads,loss] = dlfeval(@modelGradients,dlnet,dlx0,targets);
    
    % Update network 
    [dlnet,averageGrad,averageSqGrad] = adamupdate(dlnet,grads,averageGrad,averageSqGrad,iter,...
        learnRate,gradDecay,sqGradDecay);
    
    % Plot loss
    currentLoss = extractdata(loss);
    
    if plots == "training-progress"
        addpoints(lossline, iter, currentLoss);
        drawnow
    end
    
    % Plot predicted vs. real dynamics
    if mod(iter,plotFrequency) == 0
        figure(2)
        clf

        % Extract the learnt dynamics
        internalNeuralOdeLayer = dlnet.Layers(1);
        dlnetODEFcn = @(t,y) evaluateODE(internalNeuralOdeLayer, y);

        % Use ode45 to compute the solution 
        [~,y] = ode45(dlnetODEFcn, [t(1) t(end)], x0, odeOptions);
        y = y';
        
        plot(x(1,trainingTimesteps),x(2,trainingTimesteps),'r--')
        hold on
        plot(y(1,:),y(2,:),'b-')
        hold off
        D = duration(0,0,toc(start),'Format','hh:mm:ss');
        title("Iter = " + iter + ", loss = " + num2str(currentLoss) + ", Elapsed: " + string(D))
        legend('Training ground truth', 'Predicted')
    end
end

Оценка модели

Используйте обученную модель как правую сторону для той же задачи с различными начальными условиями, решите численно динамику обучения ОДУ с ode45 и сравните его с использованием истинной модели.

tPred = t;
x0Pred1 = sqrt([2;2]);
x0Pred2 = [-1;-1.5];
x0Pred3 = [0;2];
x0Pred4 = [-2;0];

[xPred1, xTrue1, err1] = predictWithOde45(dlnet,A,tPred,x0Pred1,odeOptions);
[xPred2, xTrue2, err2] = predictWithOde45(dlnet,A,tPred,x0Pred2,odeOptions);
[xPred3, xTrue3, err3] = predictWithOde45(dlnet,A,tPred,x0Pred3,odeOptions);
[xPred4, xTrue4, err4] = predictWithOde45(dlnet,A,tPred,x0Pred4,odeOptions);

Постройте график предсказанных решений относительно основной истины решений.

subplot(2,2,1)
plotTrueAndPredictedSolutions(xTrue1, xPred1, err1, "[sqrt(2) sqrt(2)]");

subplot(2,2,2)
plotTrueAndPredictedSolutions(xTrue2, xPred2, err2, "[-1 -1.5]");

subplot(2,2,3)
plotTrueAndPredictedSolutions(xTrue3, xPred3, err3, "[0 2]");

subplot(2,2,4)
plotTrueAndPredictedSolutions(xTrue4, xPred4, err4, "[-2 0]");

Функция градиентов модели

Эта функция принимает набор начальных условий, dlx0, и целевые последовательности, targets, и вычисляет потери и градиенты относительно параметров сети.

function [gradients,loss] = modelGradients(dlnet, dlX0, targets)

% Compute prediction of network
dlX = forward(dlnet,dlX0);

% Compute mean absolute error loss
loss = sum(abs(dlX - targets), 'all') / numel(dlX);

% Compute gradients
gradients = dlgradient(loss,dlnet.Learnables);

end

Эта функция создает пакет наблюдений целевой динамики.

function [dlX0, dlT] = createMiniBatch(numTimesteps, numTimesPerObs, miniBatchSize, X)
% Create batches of trajectories
s = randperm(numTimesteps - numTimesPerObs, miniBatchSize);

dlX0 = dlarray(X(:, s),'CB');
dlT = zeros([size(dlX0,1) miniBatchSize numTimesPerObs]);

for i = 1:miniBatchSize
    dlT(:, i, 1:numTimesPerObs) = X(:, s(i):(s(i) + numTimesPerObs - 1));
end
end

Функция predictWithOde45 вычисляет истинные и предсказанные решения и возвращает их вместе с соответствующей ошибкой.

function [xPred, xTrue, error] = predictWithOde45(dlnet,A,tPred,x0Pred,odeOptions)
% Use ode45 to compute the solution both with the true and the learnt
% models.

[~, xTrue] = ode45(@(t,y) A*y, tPred, x0Pred, odeOptions);

% Extract the learnt dynamics
internalNeuralOdeLayer = dlnet.Layers(1);
dlnetODEFcn = @(t,y) evaluateODE(internalNeuralOdeLayer, y);

[~,xPred] = ode45(dlnetODEFcn, tPred, x0Pred, odeOptions);
error = mean(abs(xTrue - xPred), 'all');
end

function plotTrueAndPredictedSolutions(xTrue, xPred, err, x0Str)
plot(xTrue(:,1),xTrue(:,2),'r--',xPred(:,1),xPred(:,2),'b-','LineWidth',1)
title("x_0 = " + x0Str + ", err = " + num2str(err) )
xlabel('x1')
ylabel('x2')
xlim([-2 2])
ylim([-2 2])
legend('Ground truth', 'Predicted')
end

[1] Нейронные обыкновенные дифференциальные уравнения, Рики Т. К. Чен, Юлия Рубанова, Джесси Бетанкур, Дэвид Дювено, 2019, https://arxiv.org/abs/1806.07366

[2] Численная математика, A Quarteroni, R Sacco, F Saleri - 2010

См. также

dlarray | dlfeval | dlgradient

Документация