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

 이전 발행 글에서 마저 분석하다만 Physics-Informed Neural Operator(PINO)의 코드를 분석할 예정이다. 방정식의 상세한 조건 세팅 등은 이전 글에 있으므로 검색을 통해 이번 포스팅부터 보시는 분들은 아래에 있는 링크를 통해 이전 글부터 보고 오시면 될 듯하다. 이전 글에서는 Generating Data, Fourier Neural Operator(FNO) 부분을 다뤘으며, 이번 글에서는 나머지인 Defined Loss function, Train Neural Operator 부분을 다룰 예정이다.

 

2023.10.10 - [인공지능/Deep Learning] - [PINO] Burger's Equation 예제 코드 분석 [1/2]

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

 

Defined Loss function

 Physics-Informed Neural Operator(PINO)가 Fourier Neural Operator(FNO)의 차별점은 Loss function에 물리적 정보를 포함하였다는 점이다. 이를 코드를 통해 확인해 볼 것이다. Operator $G_{\theta}$는 $\mathcal{L}_{pde}$와 $\mathcal{L} _{op}$의 합으로 인한 손실함수를 통해 훈련이 된다. 이때 $\mathcal{L}_{op}$는 다음과 같다.

 

$$\mathcal{L}_{op}(G_{\theta_i}(a),G_{\theta_o}(a)) := \Vert G_{\theta_i}(a)-G_{\theta_o}(a)\Vert^2_u$$

 

 이는 Data Loss로 일반적으로 사용하는 방식과 똑같이 취해지므로 본 글에서 굳이 코드를 붙여가며 설명하지는 않겠다. Class LpLoss가 Data Loss를 얻기 위한 Class이다. 

 

 $\mathcal{L}_{pde}$은 $\mathcal{L}_{ic}$와 $\mathcal{L}_{physics}$의 항으로 구성이 된다.

 

 $\mathcal{L}_{physics}$를 계산하기 위해 유한차분법(Finite Difference Method)을 통해 Burger's Equation을 계산한다. 유한차분법을 통해 Burger's Equation을 해석하는 방법 또한 이해하고 싶은 꼼꼼하신 분들을 위해 이와 비슷한 예제를 본문 하단에 첨부하였으니 관심 있는 사람은 확인해 보면 될 것 같다.

 

def FDM_Burgers(u, v, D=1):
    
    유한차분법을 통해 계산    
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

    Du = ut + (ux*u - v*uxx)[:,1:-1,:]
    return Du

 

 아래 코드가  $\mathcal{L}_{pde}$를 얻기 위한 코드이다. 여기서 loss_u가 $\mathcal{L}_{ic}$이며, loss_f가  $\mathcal{L}_{physics}$이다. loss_bc의 경우 주석 처리가 되어있는 모습을 확인할 수 있는데, PINO의 경우(FNO를 포함하여) 주파수 영역에서의 매핑을 취하기 때문에 경계조건을 필요로 하지 않다. 이러한 모습이 반영된 모습이다. 

 

def PINO_loss(u, u0, v):
	
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
	
    loss_u = F.mse_loss(boundary_u, u0)

    Du = FDM_Burgers(u, v)[:, :, :]
    f = torch.zeros(Du.shape, device=u.device)
    loss_f = F.mse_loss(Du, f)

    # loss_bc0 = F.mse_loss(u[:, :, 0], u[:, :, -1])
    # loss_bc1 = F.mse_loss((u[:, :, 1] - u[:, :, -1]) /
    #                       (2/(nx)), (u[:, :, 0] - u[:, :, -2])/(2/(nx)))
    return loss_u, loss_f

 

Train Neural Operator

 이제 최종 장인 모델을 훈련시키기 위한 train_2d_burger 함수를 정의한다. Github에서 제공된 코드에서는 wandb 라이브러리를 통해 모델 훈련 도중의 loss 및 여러 정보들을 확인하기 때문에 사전에 wandb 로그인을 하여 경로를 맞춰주어야 한다. train_2d_burger 함수에서 크게 봐야 될 두 가지 부분(코드 한 부분과 yaml 파일)만 살펴볼 예정이며, 아래에 있는 코드는 train_2d_burger 함수 내부의 코드의 일부이다.

 

        for x, y in train_loader:
            x, y = x.to(rank), y.to(rank)
            out = model(x).reshape(y.shape)
            data_loss = myloss(out, y)

            loss_u, loss_f = PINO_loss(out, x[:, 0, :, 0], v)
            total_loss = loss_u * ic_weight + loss_f * f_weight + data_loss * data_weight

            optimizer.zero_grad()
            total_loss.backward()
            optimizer.step()

 

 total_loss를 계산하기 위해 loss_u, loss_f, data_loss에 각각 가중치를 곱해주는 모습을 볼 수 있다. PINO 논문에서 저자는 각각의 loss에 가중치를 고려해 주면 더 좋은 성능을 얻을 수도 있다는 점을 어필하는데, 이를 보여주기 위해 각각의 loss에 가중치가 고려된 모습이다. 각 weight의 값은 yaml 파일 안에 정의되어 있는데, weight 값 외의 다양한 매개 변수들이 yaml 파일 안에 정의되어 있다.(yaml 파일 또한 같이 Github에 제공된다.) 

 

 아래 코드 블럭 내에 있는 정보들은 yaml 파일의 일부분을 가져온 것이다. 여기서 특이한 점을 한 가지 볼 수 있는데, data_loss가 0이라는 것이다. 아래의 data_loss는 코드 내에서 data_weight로 사용이 된다. 그렇다면 여기까지 잘 이해한 사람들은 엥? Data Loss를 안 쓴다고?라는 의문을 가지는 것이 당연하다. 이는 PINO가 Data가 필요 없다는 Unsupervised 하다는 특성을 보여주는 부분으로 data_loss에 다른 숫자를 넣어 실제로 사용을 하여도 결과에는 지장이 없다.

 

  ic_loss: 10.0
  f_loss: 1.0
  data_loss: 0.0

log:
  project: PINO-burgers-pretrain
  group: xxxxxxmmdd
  entity: xxxxxxx132

 

 또한 위의 yaml 파일의 log가 wandb 경로를 잡아주는 부분으로 자신의 wandb의 project와 group, entity에 맞추어 재설정되면 모델이 훈련 중일 때 wandb에서 실시간으로 훈련되는 모습을 확인할 수 있다.

 

글을 마치며

 Neural Operator를 Test 하는 코드는 Train과 비슷하니 해당 부분까지 잘 이해하신 분들이라면 큰 어려움이 없을 거라 생각하고 생략하도록 한다. PINO의 Burger's Equation 코드의 경우 구글 코랩 환경에서도 충분히 잘 돌아가니, 한번 실행해 봐도 좋을 듯하다. 마지막으로 논문 링크와 코드가 있는 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

유한차분법 Burger's Equation 예제 : https://github.com/nishantsule/Burgers-Equation/blob/master/BurgersEq_1D.ipynb