[PINO] Burger's Equation 예제 코드 분석 [1/2]

 이전의 PINN에 이어 Physics-Informed Neural Operator(PINO)의 알고리즘 작동 원리 및 코드 짜임새를 살펴보기 위해 Burger's Equation의 예제 코드를 분석해 볼 예정이다. Burger's Equation은 식 구성이 단순해선지는 모르겠지만, 여러 프레임워크가 시작 예제 느낌으로 코드를 많이 제공하기 때문에, 가능한 한 앞으로도 딥러닝 프레임워크의 코드 분석에 있어서는 Burger's Equation을 사용하지 않을까 싶다. 동일한 방정식을 토대로 한 여러 코드를 보다 보면 방정식을 확실하게 알고 있다 보니, 그만큼 각 프레임워크의 차이도 구별하기 쉬운데, 이 글을 읽는 여러분들이 이런 효과를 조금이라도 느끼셨으면 하는 마음도 있다.

 

 이번 포스팅에 사용할 코드는 "Physics-Informed Neural Operator for Learning Partial Differential Equations"에서 배포한 코드이며 그 중 Burger's Equation을 분석할 것이다. 맨 아래에 논문 링크와 Github 링크가 나와있으니, 필요할 경우 적절하게 참고하면 된다. PINO에 대해 간단한 설명은 이전 글에 적어두었으니 참고하면 될 듯하다. 또한 Github 상에 찢겨 있는 코드들을 필자가 설명을 위해 편의상 중간중간 끼워 넣었으니 코드를 직접 실행시켜 보며 비교해 보는 경우 참고하면 될 듯하다. 또한 코드를 다 포스팅에 올리자니 너무 가독성이 떨어져 설명에 직접적으로 필요한 부분을 제외하고는 "~~~"로 잘랐으니 참고하면 될 듯하다.

 

2023.09.19 - [인공지능/Deep Learning] - [FNO] Fourier Neural Operator에 대해

[FNO] Fourier Neural Operator에 대해

 

Problem Setting

 Burger's Equation의 경우 논문의 Section 4.1을 구현한 것으로, 간단하게 문제에 대해 살펴보도록 하자.

 

$$\partial_tu(x,t)+\partial_x(u^2(x,t)/2)=\nu\partial_{xx}u(x,t),\quad x ∈ (0,1),t∈(0,1]$$

$$u(x,0)=u_0(x), \quad x∈(0,1)$$

 

 위 방정식은 $u_0∈L^2_{per}((0,1); R)$의 주기적 경계 조건을 가지고 있다. $\nu$는 점도 계수로 해당 예제에서는 0.01 수치가 사용된다. 본 예제 코드에서는 초기 조건을 솔루션 $G_† : u_0 → u|_{[0,1]}.$에 매핑하는 연산자를 학습하는 것을 목표로 한다.

 

Generating Data

 먼저 Neural Operator를 훈련시키기 위해 훈련 데이터인 Matlab 데이터 파일을 불러올 수 있는 클래스를 정의한다. 본 예제에서는 MatReader와 BurgersLoader를 사용하는데, MatReader에서는 일반적인 field를 읽거나, gpu 사용 여부 등의 기본적인 함수들이 들어가 있다. 

 

class BurgersLoader(object):
    def __init__(self, datapath, nx=2 ** 10, nt=100, sub=8, sub_t=1, new=False):
        dataloader = MatReader(datapath)
        self.sub = sub
        self.sub_t = sub_t
        self.s = nx // sub
        self.T = nt // sub_t
        self.new = new
        if new:
            self.T += 1
        self.x_data = dataloader.read_field('input')[:, ::sub]
        self.y_data = dataloader.read_field('output')[:, ::sub_t, ::sub]
        self.v = dataloader.read_field('visc').item()

 

 BurgersLoader의 생성자 init이 input으로 받는 매개변수들의 의미는 다음과 같다.

datapath : Matlab data가 위치해 있는 경로

nx, nt : 공간, 시간 그리드의 총 포인트 수

sub, sub_t : 공간, 시간 단위의 하위 샘플링 비율

 

 위의 코드를 확인해 보면 $s$와 $T$라는 변수에 그리드의 총 포인트 수에 하위 샘플링 비율을 나눠 새로운 공간과 시간 단위의 변수를 저장해 주는 모습을 보인다. 이 두 변수의 쓰임은 아래의 make_loader에서 확인할 수 있다. 아래의 코드에서 gridx와 gridt가 그리드의 총 포인트 수인 nx, nt가 아닌, 이 값에 샘플링 비율을 나눈 $s$와 $T$를 기준으로 만들어진다. FNO와 PINO의 경우는 데이터의 Resolution과 관련 없이 evaluate를 할 수 있는데, 그 이유를 여기에서 확인할 수 있다.

 

    def make_loader(self, n_sample, batch_size, start=0, train=True):
            Xs = self.x_data[start:start + n_sample]
            ys = self.y_data[start:start + n_sample]

            if self.new:
                gridx = torch.tensor(np.linspace(0, 1, self.s + 1)[:-1], dtype=torch.float)
                gridt = torch.tensor(np.linspace(0, 1, self.T), dtype=torch.float)
            else:
                gridx = torch.tensor(np.linspace(0, 1, self.s), dtype=torch.float)
                gridt = torch.tensor(np.linspace(0, 1, self.T + 1)[1:], dtype=torch.float)
            gridx = gridx.reshape(1, 1, self.s)
            gridt = gridt.reshape(1, self.T, 1)

            Xs = Xs.reshape(n_sample, 1, self.s).repeat([1, self.T, 1])
            Xs = torch.stack([Xs, gridx.repeat([n_sample, self.T, 1]), gridt.repeat([n_sample, 1, self.s])], dim=3)
            dataset = torch.utils.data.TensorDataset(Xs, ys)
            if train:
                loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=True)
            else:
                loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=False)
            return loader

 

Fourier Neural Operator(FNO)

 이전 FNO 글에서 언급한 것처럼 FNO는 Fourier Transform, Linear Transform on the Lower Fouier Modes, Inverse Fourier Transform을 거친 값과 Bias term을 거친 값이 더해지며 활성화 함수를 거친다. PINO 또한 물리적 정보를 손실함수로 사용하는 것이기 때문에 손실 함수의 차이만 있으며, 기본적으로는 FNO와 동일한 네트워크를 사용한다. 

 

 Fouier Layer에서는 mode가 중요한 역할을 하는데, 아래 코드에서도 mode1, mode2를 인수로 받는 모습을 확인할 수 있다. 이는 각 첫 번째 차원과 두 번째 차원에서 사용할 Fourier mode의 개수를 나타내는 인자로, Fourier Transform을 수행할 때, 입력 데이터의 주파수 성분들을 나타내는 일종의 기저함수의 역할을 한다.

 

import torch.nn as nn
#from models.basics import SpectralConv2d
from models.utils import _get_act, add_padding2, remove_padding2


class FNO2d(nn.Module):
    def __init__(self, modes1, modes2,
                 width=64, fc_dim=128,
                 layers=None,
                 in_dim=3, out_dim=1,
                 act='gelu',
                 pad_ratio=[0., 0.]):
        super(FNO2d, self).__init__()
        """
        Args:
            - modes1: list of int, number of modes in first dimension in each layer
            - modes2: list of int, number of modes in second dimension in each layer
            - width: int, optional, if layers is None, it will be initialized as [width] * [len(modes1) + 1]
            - in_dim: number of input channels
            - out_dim: number of output channels
            - act: activation function, {tanh, gelu, relu, leaky_relu}, default: gelu
            - pad_ratio: list of float, or float; portion of domain to be extended. If float, paddings are added to the right.
            If list, paddings are added to both sides. pad_ratio[0] pads left, pad_ratio[1] pads right.
        """

		~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

        self.fc0 = nn.Linear(in_dim, layers[0])

        self.sp_convs = nn.ModuleList([SpectralConv2d(
            in_size, out_size, mode1_num, mode2_num)
            for in_size, out_size, mode1_num, mode2_num
            in zip(self.layers, self.layers[1:], self.modes1, self.modes2)])

        self.ws = nn.ModuleList([nn.Conv1d(in_size, out_size, 1)
                                 for in_size, out_size in zip(self.layers, self.layers[1:])])

        self.fc1 = nn.Linear(layers[-1], fc_dim)
        self.fc2 = nn.Linear(fc_dim, layers[-1])
        self.fc3 = nn.Linear(layers[-1], out_dim)
        self.act = _get_act(act)

 

 위의 코드에서 sp_convs(SpectralConv2d), ws(Conv1d)가 정의된 모습을 확인할 수 있다. SpecralConv2d Layer의 경우 Fourier Transform을 수행하는 Layer로 Fourier Layer의 3단계(그림의 노란 박스) 부분을 수행한다. 또한 Conv1d Layer는 1차원 Convolution 연산을 수행하여 위 그림의 Bias Term을 얻을 수 있도록 한다. 이를 직접적으로 아래 코드에서 x1에 SpectralConv2d, x2에 Conv1d를 적용하여 둘을 더해 FNO를 적용시켜 주는 모습을 확인할 수 있다.

 

 또한 Fourier Layer를 거치기 전에는 항상 Lifting Operator를 거친 후 Channel을 올려주고 최종적으로 Fourier Layer를 다 거치고 나서는 Projection Operator로 Channel을 낮춰주는데, 아래 코드에서 permute를 통해 Lifting Operator와 Projection Operator의 역할을 취한다.

 

 그렇다면 왜 첫 번째 layer에서만 FNO를 적용시키고 fc1, fc2, fc3등은 그냥 Linear Transform을 거치는지 의문을 가질 수 있다. 이는 논문에서 저자가 PINO에서 FNO와 Linear Transform을 적절히 조합하면 보다 더 좋은 성능을 얻어낼 수 있다고 서술하는데, 그렇기 때문에 해당 예제에서도 둘을 섞어 쓰는 것이 아닌가 싶다.

 

    def forward(self, x):

        ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

        x = self.fc0(x)
        x = x.permute(0, 3, 1, 2)   # B, C, X, Y

        x = add_padding2(x, num_pad1, num_pad2)
        size_x, size_y = x.shape[-2], x.shape[-1]

        for i, (speconv, w) in enumerate(zip(self.sp_convs, self.ws)):
            x1 = speconv(x)
            x2 = w(x.view(batchsize, self.layers[i], -1)).view(batchsize, self.layers[i+1], size_x, size_y)
            x = x1 + x2
            if i != length - 1:
                x = self.act(x)

        x = remove_padding2(x, num_pad1, num_pad2)
        x = x.permute(0, 2, 3, 1)  # B, X, Y, C

        x = self.fc1(x)
        x = self.act(x)
        x = self.fc2(x)
        x = self.act(x)
        x = self.fc3(x)
        return x

 

 이제 가장 중요한 실제로 SpectralConv2d에서 Fourier Layer가 어떻게 작동되는지를 봐야 할 것이다. compl_mul2d는 행렬곱을 수행하는 함수이며, SpectralConv2d의 init 생성자 부분에서는 weights1과 weighs2를 정의함으로써, Fourier Layer의 노란 박스 부분에서도 Optimizing 될 수 있게 nn.Parameter로 학습 가능한 매개변수로 설정해 준다. 

 

 forward 코드를 보면 rfftn을 거쳐 Fourier Transform을 하는 모습과 irttn을 통해 Inverse Fourier Transform을 거치는 모습을 한눈에 확인할 수 있다. #Multiply relevant Fourier modes의 부분이 High Frequency mode를 제외하고 Linear Transform을 거치는 부분이다. modes1을 기준으로 그 안의 범위만을 살리며, 그 밖의 부분은 버리는 모습을 확인할 수 있다. 그렇다면 왜 modes2는 고정적으로 두고 사용하냐? 그 이유는 정확하게 이해는 하지 못하였지만, 유튜브 강의들을 몇 개 확인해 보니 대칭으로 인한 conjugacy 한 특성 때문이란다.. 어렴풋이 감이 잡히는 부분은 있지만, 글로 풀 정도로 정확하게 이해하지는 못하였다.

def compl_mul2d(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
	~~~~~~~~~~~~~~~~~~~~~~
    return res

class SpectralConv2d(nn.Module):
	def __init__(self, in_channels, out_channels, modes1, modes2):
		~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
       	self.weights1 = nn.Parameter(
            self.scale * torch.rand(in_channels, out_channels, self.modes1, self.modes2, dtype=torch.cfloat))
        self.weights2 = nn.Parameter(
            self.scale * torch.rand(in_channels, out_channels, self.modes1, self.modes2, dtype=torch.cfloat))

    def forward(self, x):
		~~~~~~~~~~~~~~~~~~~~~~~~~~~
        x_ft = torch.fft.rfftn(x, dim=[2, 3])

        # Multiply relevant Fourier modes
        out_ft = torch.zeros(batchsize, self.out_channels, x.size(-2), x.size(-1) // 2 + 1, device=x.device,
                                dtype=torch.cfloat)
        out_ft[:, :, :self.modes1, :self.modes2] = \
            compl_mul2d(x_ft[:, :, :self.modes1, :self.modes2], self.weights1)
        out_ft[:, :, -self.modes1:, :self.modes2] = \
            compl_mul2d(x_ft[:, :, -self.modes1:, :self.modes2], self.weights2)

        # Return to physical space
        x = torch.fft.irfftn(out_ft, s=(x.size(-2), x.size(-1)), dim=[2, 3])
        return x

 

글을 마치며

 이번 글에서는 PINO의 Generating Data의 부분과 Fourier Neural Operator(FNO) 부분을 확인해 보았다. 다음 글에서는 PINO가 FNO와 구별되는 주 이유인 Loss Function에 대한 부분과 최종적으로 어떻게 Train 하는지 확인해 볼 예정이다. Fourier Neural Operator를 검색하고 오신 분이더라도 이번 글의 구성은 FNO와 PINO과 차이가 없으므로 한번 읽어봐도 좋을 듯하다. 마지막으로 PINO의 논문 링크와 Github 링크를 남겨두고 글을 마치도록 하겠다.

 

논문 링크 : https://doi.org/10.48550/arXiv.2111.03794

Physics-Informed Neural Operator for Learning Partial Differential Equations

Github 링크 : https://github.com/neuraloperator/physics_informed.git

GitHub - neuraloperator/physics_informed