深層学習を使用した偏微分方程式の求解

この例では、深層学習を使用してバーガース方程式を解く方法を示します。

バーガース方程式は、応用数学のさまざまな分野に登場する偏微分方程式 (PDE) です。特に、流体力学、非線形音響学、気体力学、および交通流の分野です。

与えられた計算領域 $[- 1, 1] \times [0, 1]$ について、この例では、Physics Informed Neural Network (PINN) [1] を使用して、サンプル $(x, t)$ を入力として受け取る多層パーセプトロンニューラルネットワークに学習させます。ここで、 $x \in [- 1, 1]$ は空間変数、 $t \in [0, 1]$ は時間変数です。また、 $u (x, t)$ (u は次のバーガース方程式の解) が返されます。

$\frac{\partial u}{\partial t} + u \frac{\partial u}{\partial x} - \frac{0.01}{π} \frac{\partial^{2} u}{\partial x^{2}} = 0,$

初期状態は $u (x, t = 0) = - s i n (π x)$ 、境界条件は $u (x = - 1, t) = 0$ および $u (x = 1, t) = 0$ とします。

この例では、与えられた入力 $(x, t)$ について、ネットワークの出力 $u (x, t)$ がバーガース方程式、境界条件、および初期状態を満たすように強制することで、モデルに学習させます。

このモデルの学習には、事前のデータ収集を必要としません。PDE の定義および制約を使用して、データを生成できます。

学習データの生成

モデルに学習させるには、境界条件を強制し、初期状態を強制し、バーガース方程式を満たす選点のデータセットが必要です。

25 個の等間隔の時間点を選択して、各境界条件 $u (x = - 1, t) = 0$ および $u (x = 1, t) = 0$ を強制します。

numBoundaryConditionPoints = [25 25];

x0BC1 = -1*ones(1,numBoundaryConditionPoints(1));
x0BC2 = ones(1,numBoundaryConditionPoints(2));

t0BC1 = linspace(0,1,numBoundaryConditionPoints(1));
t0BC2 = linspace(0,1,numBoundaryConditionPoints(2));

u0BC1 = zeros(1,numBoundaryConditionPoints(1));
u0BC2 = zeros(1,numBoundaryConditionPoints(2));

50 個の等間隔の時間点を選択して、初期状態 $u (x, t = 0) = - s i n (π x)$ を強制します。

numInitialConditionPoints  = 50;

x0IC = linspace(-1,1,numInitialConditionPoints);
t0IC = zeros(1,numInitialConditionPoints);
u0IC = -sin(pi*x0IC);

初期状態および境界条件について、データをグループ化します。

X0 = [x0IC x0BC1 x0BC2];
T0 = [t0IC t0BC1 t0BC2];
U0 = [u0IC u0BC1 u0BC2];

10,000 個の点を選択し、ネットワークの出力がバーガース方程式を満たすように強制します。

numInternalCollocationPoints = 10000;

pointSet = sobolset(2);
points = net(pointSet,numInternalCollocationPoints);

dataX = 2*points(:,1)-1;
dataT = points(:,2);

学習データを格納する配列データストアを作成します。

ds = arrayDatastore([dataX dataT]);

深層学習モデルの定義

20 個の隠れニューロンがある 9 つの全結合演算で多層パーセプトロンアーキテクチャを定義します。最初の全結合演算には、入力 $x$ および $t$ に対応する 2 つの入力チャネルがあります。最後の全結合演算には、1 つの出力 $u (x, t)$ があります。

モデルパラメーターの定義と初期化

各演算のパラメーターを定義して struct に含めます。parameters.OperationName.ParameterName の形式を使用します。ここで、parameters は struct、OperationName は演算名 ("fc1" など)、ParameterName はパラメーター名 ("Weights" など) です。

層の数および各層のニューロンの数を指定します。

numLayers = 9;
numNeurons = 20;

最初の全結合演算のパラメーターを初期化します。最初の全結合演算には 2 つの入力チャネルがあります。

parameters = struct;

sz = [numNeurons 2];
parameters.fc1.Weights = initializeHe(sz,2);
parameters.fc1.Bias = initializeZeros([numNeurons 1]);

残りの中間の各全結合演算のパラメーターを初期化します。

for layerNumber=2:numLayers-1
    name = "fc"+layerNumber;

    sz = [numNeurons numNeurons];
    numIn = numNeurons;
    parameters.(name).Weights = initializeHe(sz,numIn);
    parameters.(name).Bias = initializeZeros([numNeurons 1]);
end

最後の全結合演算のパラメーターを初期化します。最後の全結合演算には、1 つの出力チャネルがあります。

sz = [1 numNeurons];
numIn = numNeurons;
parameters.("fc" + numLayers).Weights = initializeHe(sz,numIn);
parameters.("fc" + numLayers).Bias = initializeZeros([1 1]);

ネットワークパラメーターを表示します。

parameters

parameters = struct with fields:
    fc1: [1×1 struct]
    fc2: [1×1 struct]
    fc3: [1×1 struct]
    fc4: [1×1 struct]
    fc5: [1×1 struct]
    fc6: [1×1 struct]
    fc7: [1×1 struct]
    fc8: [1×1 struct]
    fc9: [1×1 struct]

最初の全結合層のパラメーターを表示します。

parameters.fc1

ans = struct with fields:
    Weights: [20×2 dlarray]
       Bias: [20×1 dlarray]

モデルとモデル損失関数の定義

この例の最後のモデル関数の節にリストされている関数 model を作成します。この関数は深層学習モデルの出力を計算します。関数 model は、モデルパラメーターとネットワーク入力を入力として受け取り、モデル出力を返します。

この例の最後のモデル損失関数の節にリストされている関数 modelLoss を作成します。この関数は、モデルパラメーター、ネットワーク入力、初期状態、および境界条件を入力として受け取り、学習可能なパラメーターについての損失および損失の勾配を返します。

学習オプションの指定

ミニバッチサイズを 1000 としてモデルに 3000 エポック学習させます。

numEpochs = 3000;
miniBatchSize = 1000;

GPU が利用できる場合に GPU で学習を行うには、実行環境を "auto" に指定します。GPU を使用するには、Parallel Computing Toolbox™ とサポートされている GPU デバイスが必要です。サポートされているデバイスについては、GPU 計算の要件 (Parallel Computing Toolbox) (Parallel Computing Toolbox) を参照してください。

executionEnvironment = "auto";

ADAM 最適化オプションを指定します。

initialLearnRate = 0.01;
decayRate = 0.005;

ネットワークの学習

カスタム学習ループを使用してネットワークに学習させます。

学習中にデータのミニバッチを処理および管理する minibatchqueue オブジェクトを作成します。各ミニバッチで次を行います。

データを、次元ラベル 'BC' (batch、channel) で形式を整えます。既定では、minibatchqueue オブジェクトは、基となる型が single の dlarray オブジェクトにデータを変換します。
変数 executionEnvironment の値に従って、GPU で学習を行います。既定では、minibatchqueue オブジェクトは、GPU が利用可能な場合、各出力を gpuArray に変換します。

mbq = minibatchqueue(ds, ...
    MiniBatchSize=miniBatchSize, ...
    MiniBatchFormat="BC", ...
    OutputEnvironment=executionEnvironment);

初期状態と境界条件を dlarray に変換します。入力データ点について、次元 "CB" (channel、batch) で形式を整えます。

X0 = dlarray(X0,"CB");
T0 = dlarray(T0,"CB");
U0 = dlarray(U0);

学習に GPU を使用する場合は、初期条件と境界条件を gpuArray に変換します。

if (executionEnvironment == "auto" && canUseGPU) || (executionEnvironment == "gpu")
    X0 = gpuArray(X0);
    T0 = gpuArray(T0);
    U0 = gpuArray(U0);
end

Adam ソルバーのパラメーターを初期化します。

averageGrad = [];
averageSqGrad = [];

関数 dlaccelerate を使用して、モデル損失関数を高速化します。詳細については、Accelerate Custom Training Loop Functionsを参照してください。

accfun = dlaccelerate(@modelLoss);

学習の進行状況プロットを初期化します。

figure
C = colororder;
lineLoss = animatedline(Color=C(2,:));
ylim([0 inf])
xlabel("Iteration")
ylabel("Loss")
grid on

ネットワークに学習をさせます。

それぞれの反復で次を行います。

ミニバッチキューからデータのミニバッチを読み取ります。
高速化したモデル損失関数および関数 dlfeval を使用して、モデルの損失と勾配を評価します。
学習率を更新します。
関数 adamupdate を使用して学習可能なパラメーターを更新します。

各エポックの最後に、損失値で学習プロットを更新します。

start = tic;

iteration = 0;

for epoch = 1:numEpochs
    reset(mbq);

    while hasdata(mbq)
        iteration = iteration + 1;

        XT = next(mbq);
        X = XT(1,:);
        T = XT(2,:);

        % Evaluate the model loss and gradients using dlfeval and the
        % modelLoss function.
        [loss,gradients] = dlfeval(accfun,parameters,X,T,X0,T0,U0);

        % Update learning rate.
        learningRate = initialLearnRate / (1+decayRate*iteration);

        % Update the network parameters using the adamupdate function.
        [parameters,averageGrad,averageSqGrad] = adamupdate(parameters,gradients,averageGrad, ...
            averageSqGrad,iteration,learningRate);
    end

    % Plot training progress.
    loss = double(gather(extractdata(loss)));
    addpoints(lineLoss,iteration, loss);

    D = duration(0,0,toc(start),Format="hh:mm:ss");
    title("Epoch: " + epoch + ", Elapsed: " + string(D) + ", Loss: " + loss)
    drawnow
end

ヒット率と占有率をチェックすることにより、高速化した関数の有効性をチェックします。

accfun

accfun = 
  AcceleratedFunction with properties:

          Function: @modelLoss
           Enabled: 1
         CacheSize: 50
           HitRate: 99.9967
         Occupancy: 2
         CheckMode: 'none'
    CheckTolerance: 1.0000e-04

モデルの精度の評価

$t$ の値が 0.25、0.5、0.75、および 1 である場合について、 $l^{2}$ 誤差を使用して、深層学習モデルの予測値とバーガース方程式の真の解を比較します。

モデルをテストする際のターゲット時間を設定します。各時間について、[-1,1] の範囲内の 1001 個の等間隔の点で解を計算します。

tTest = [0.25 0.5 0.75 1];
numPredictions = 1001;
XTest = linspace(-1,1,numPredictions);

figure

for i=1:numel(tTest)
    t = tTest(i);
    TTest = t*ones(1,numPredictions);

    % Make predictions.
    XTest = dlarray(XTest,"CB");
    TTest = dlarray(TTest,"CB");
    UPred = model(parameters,XTest,TTest);

    % Calculate true values.
    UTest = solveBurgers(extractdata(XTest),t,0.01/pi);

    % Calculate error.
    err = norm(extractdata(UPred) - UTest) / norm(UTest);

    % Plot predictions.
    subplot(2,2,i)
    plot(XTest,extractdata(UPred),"-",LineWidth=2);
    ylim([-1.1, 1.1])

    % Plot true values.
    hold on
    plot(XTest, UTest, "--",LineWidth=2)
    hold off

    title("t = " + t + ", Error = " + gather(err));
end

subplot(2,2,2)
legend("Predicted","True")

プロットは、予測が真の値にどれだけ近いかを示しています。

バーガース方程式の求解関数

関数 solveBurgers は、[2] で概説されているように、時間 t におけるバーガース方程式の真の解を返します。

function U = solveBurgers(X,t,nu)

% Define functions.
f = @(y) exp(-cos(pi*y)/(2*pi*nu));
g = @(y) exp(-(y.^2)/(4*nu*t));

% Initialize solutions.
U = zeros(size(X));

% Loop over x values.
for i = 1:numel(X)
    x = X(i);

    % Calculate the solutions using the integral function. The boundary
    % conditions in x = -1 and x = 1 are known, so leave 0 as they are
    % given by initialization of U.
    if abs(x) ~= 1
        fun = @(eta) sin(pi*(x-eta)) .* f(x-eta) .* g(eta);
        uxt = -integral(fun,-inf,inf);
        fun = @(eta) f(x-eta) .* g(eta);
        U(i) = uxt / integral(fun,-inf,inf);
    end
end

end

モデル損失関数

与えられた入力 $(x, t)$ について、ネットワークの出力 $u (x, t)$ がバーガース方程式、境界条件、および初期状態を満たすように強制することで、モデルに学習させます。特に、次の 2 つの量が損失の最小化に寄与します。

$loss = {MSE}_{f} + {MSE}_{u}$ ,

ここで、 ${MSE}_{f} = \frac{1}{N_{f}} \sum_{i = 1}^{N_{f}} {| f (x_{f}^{i}, t_{f}^{i}) |}^{2}$ および ${MSE}_{u} = \frac{1}{N_{u}} \sum_{i = 1}^{N_{u}} {| u (x_{u}^{i}, t_{u}^{i}) - u^{i} |}^{2}$ です。

ここで、 ${x_{u}^{i}, t_{u}^{i}}_{i = 1}^{N_{u}}$ は、計算領域の境界上の選点に対応し、境界条件と初期状態の両方を計上します。 ${x_{f}^{i}, t_{f}^{i}}_{i = 1}^{N_{f}}$ は、領域内部の点です。

${MSE}_{f}$ を計算するには、モデルの出力 $u$ の微分 $\frac{\partial u}{\partial t}, \frac{\partial u}{\partial x}, \frac{\partial^{2} u}{\partial x^{2}}$ を求める必要があります。

関数 modelLoss は、モデルパラメーター parameters、ネットワーク入力 X および T、初期状態および境界条件である X0、T0、および U0 を入力として受け取り、学習可能なパラメーターについての損失および損失の勾配を返します。

function [loss,gradients] = modelLoss(parameters,X,T,X0,T0,U0)

% Make predictions with the initial conditions.
U = model(parameters,X,T);

% Calculate derivatives with respect to X and T.
gradientsU = dlgradient(sum(U,"all"),{X,T},EnableHigherDerivatives=true);
Ux = gradientsU{1};
Ut = gradientsU{2};

% Calculate second-order derivatives with respect to X.
Uxx = dlgradient(sum(Ux,"all"),X,EnableHigherDerivatives=true);

% Calculate lossF. Enforce Burger's equation.
f = Ut + U.*Ux - (0.01./pi).*Uxx;
zeroTarget = zeros(size(f), "like", f);
lossF = mse(f, zeroTarget);

% Calculate lossU. Enforce initial and boundary conditions.
U0Pred = model(parameters,X0,T0);
lossU = mse(U0Pred, U0);

% Combine losses.
loss = lossF + lossU;

% Calculate gradients with respect to the learnable parameters.
gradients = dlgradient(loss,parameters);

end

モデル関数

この例で学習させるモデルは、各演算の間に tanh 演算を行う一連の全結合演算で構成されています。

モデル関数は、モデルパラメーター parameters とネットワーク入力 X および T を入力として受け取り、モデル出力 U を返します。

function U = model(parameters,X,T)

XT = [X;T];
numLayers = numel(fieldnames(parameters));

% First fully connect operation.
weights = parameters.fc1.Weights;
bias = parameters.fc1.Bias;
U = fullyconnect(XT,weights,bias);

% tanh and fully connect operations for remaining layers.
for i=2:numLayers
    name = "fc" + i;

    U = tanh(U);

    weights = parameters.(name).Weights;
    bias = parameters.(name).Bias;
    U = fullyconnect(U, weights, bias);
end

end

参考文献

Maziar Raissi, Paris Perdikaris, and George Em Karniadakis, Physics Informed Deep Learning (Part I): Data-driven Solutions of Nonlinear Partial Differential Equations https://arxiv.org/abs/1711.10561
C. Basdevant, M. Deville, P. Haldenwang, J. Lacroix, J. Ouazzani, R. Peyret, P. Orlandi, A. Patera, Spectral and finite difference solutions of the Burgers equation, Computers & fluids 14 (1986) 23–41.