뛰는 놈 위에 나는 공대생

Neural ODE를 구현해놓은 코드는 torchdiffeq인데 학습이 너무 느리다는 생각이 들었다. 여러가지를 테스트해봐야 하는 입장에서 아무리 좋은 GPU를 써도 코드가 뒷받침되지 않으면 학습하는 데 시간이 오래 걸린다. 최근 JAX가 이런 측면에서 효과적이라는 것을 알아서 JAX 기반의 Neural ODE 코드를 찾아보았다. https://docs.kidger.site/diffrax/ Diffrax Diffrax in a nutshell Diffrax is a JAX-based library providing numerical differential equation solvers. Features include: ODE/SDE/CDE (ordinary/stochastic/controlled) sol..

기존에 자주 사용되면 파이토치나 텐서플로에서도 그렇듯이 학습한 모델을 저장하는 것은 필수이다. 나중에 다시 결과를 출력해야하거나 Transfer learning 등에 활용해야하기 때문이다. 이 글에서는 JAX 모델을 학습한 다음, 저장하고 다시 로드하는 방법에 대해서 다룬다. 1. 학습 후 저장 이 글에서는 model이라는 class 안에 optimizers를 정의하고 그 안에 있는 loss 등의 함수로 학습을 하고 있었다. class 안에 self.opt_init, \ self.opt_update, \ self.get_params = optimizers.adam(optimizers.exponential_decay(1e-3, decay_steps=1000, decay_rate=0.9)) self.opt_..

JAX를 통해 병렬로 뉴럴 네트워크를 학습하는 예제를 고민하였다.JAX에서 제공해주는 예제도 있지만 이는 아주 심플한 선형 모델의 파라미터를 regression하는 문제이기 때문에 실제 뉴럴 네트워크 모델과는 괴리가 좀 있어서 직접 예제를 만들었다. JAX 0.3.1. 버전 1. JAX에서 제공하는 예제 import jaxjax.devices()>> [GpuDevice(id=0, process_index=0), GpuDevice(id=1, process_index=0), GpuDevice(id=2, process_index=0), GpuDevice(id=3, process_index=0)] 필자는 gpu 4개를 가지고 병렬 컴퓨팅을 사용했다. import numpy as npimport jax.numpy..

머신러닝, 딥러닝을 공부하다보면 다양한 에러에 대해서 배우게 된다.이 글에서는 이러한 용어들의 혼동을 막고자 종합적으로 정리하는 글이다. 1. Model bias, Estimation bias, Estimation variance  그 중에 하나가 model bias, estimation bias, estimation variance이다.    위 그림을 보면 이 개념을 설명하는 결과라고 볼 수 있다. Model bias는 실제 값과 내가 모델로 만든 값 중에서 가장 실제값을 가깝게 반영하는 모델에서 발생하는 바이어스이다. 즉, 내가 정한 모델은 아무리 최적화를 시켜도 truth 값과 완벽하게 일치하지 않을 수 있다. 예를 들어 비선형 모델인데 내가 아무리 선형 모델로 fitting을 하려고 해도 모델 ..

JAX를 쓰다가 너무 많은 양의 데이터를 쓰다보니 메모리 부족(OOM: Out of memory) 현상을 겪었다. 근본적인 해결책은 달리 없다. 데이터가 너무 많아서 생기는 문제이니 데이터 양을 줄이던가 아니면 병렬 컴퓨팅을 하는 방법이 있다. 병렬 컴퓨팅을 간단한 코드에서는 실행해보았는데 큰 네트워크에서는 해본 적이 없다. 일단 임시방편으로는 다음과 같다. 메모리 부족이 쉽게 발생할 수 있는 이유는 JAX에서 처음에 import를 할 때 대부분의 메모리를 미리 할당해놓기 때문이다. 따라서 이 preallocation을 막거나 줄이면 도움이 된다. 1. Preallocation 중단 XLA_PYTHON_CLIENT_PREALLOCATE=false # 구체적으로는 다음과 같이 구현한다. import os ..

이번에 리눅스 환경에서 CUDA, cuDNN을 설치하면서 있었던 시행착오를 기록하는 글이다. 기본적으로 Linux 환경에서 설치하는 방법은 매뉴얼에 잘 나와있어서 이 링크를 참고하면 되기는 하는데 디테일하게는 고민할 부분들이 있다. 이 글에서 주의할 점은 1) 나는 이미 CUDA, cuDNN을 다른 사람이 설치해놓은 버전이 있었다. (그 버전들은 root에 설치되어 있었다.) 2) 추가적으로 CUDA 다른 버전을 쓰고 싶어서 설치하기 시작했다. 3) 우분투 user 중에 나의 계정이 있지만 내 계정에 local하게 설치하는 것이 아니라 모든 유저가 쓸 수 있도록 설치하였다. 나만 쓰고 싶으면 내 home directory에서 시작되는 경로에 설치하면 된다. 공용 서버컴의 경우에는 root 계정에 이미 설..

JAX에서 사용 가능한 device를 찾는 방법은 다음과 같다. import jax jax.devices() >> [GpuDevice(id=0, process_index=0), GpuDevice(id=1, process_index=0), GpuDevice(id=2, process_index=0), GpuDevice(id=3, process_index=0)] 여기에서는 device가 4개이기 때문에 하나를 지정해서 쓰고 싶을 수 있다. 방법을 찾아보니 document에서는 jax.default_device = jax.devices("gpu")[2] # default로 세번째 gpu를 쓰고 싶은 경우 다음과 같이 쓰면 된다고 했지만 실제로는 적용이 되지 않았다. 좀 더 확실한 방법으로는, JAX를 impor..

이 글은 JAX 버전 맞추느라 여러 🐶고생한 경험을 바탕으로 작성하였다. 0. 요구 버전에 대한 이해 JAX는 설치할 때 요구하는 버전이 있다. 개별 gpu에 따라도 달라져서 까다롭긴한데 JAX currently ships three CUDA wheel variants: CUDA 12.0 and CuDNN 8.8. CUDA 11.8 and CuDNN 8.6. CUDA 11.4 and CuDNN 8.2. This wheel is deprecated and will be discontinued with jax 0.4.8. 위의 세 버전이 가능하다고 하는데, 이는 최신 JAX 버전(230412 기준)에 따른 것이다. CUDA 11.4라고 적혀있는 경우에는 11.4 이상이면서 CuDNN 8.2 이상이면 된다. ..

JAX로 학습하는 도중에 NAN값이 나와서 어디부터 원인인지 찾기가 어려웠다. 이럴 때는 아래 코드를 추가하면 된다. from jax.config import config config.update("jax_debug_nans", True) 이렇게 할 경우에 NAN이 발생하는 즉시 어떤 코드에서 문제가 발생하는지를 알려주고 코드가 종료된다.

Gaussian process에서 사용하는 커널 종류는 다양할 수 있지만 여기서는 Radial Basis Fuction을 이용해서 gaussian process 샘플들을 구하고 이에 대한 관찰을 시각화하는 방법에 대해서 이야기한다. RBF 함수는 Paris Perdikaris 교수님의 수업자료를 참고하였다. $k(x_1,x_2)=\eta \exp\left( \dfrac{(x_1 -x_2)^{2}}{2l^2}\right)$ 커널함수가 이렇게 설정되어 있을 때 우리가 조절할 수 있는 파라미터는 scale factor인 $\eta$와 length인 $l$이다. 개념적으로 생각하였을 때 random process인 gaussian process는 $\mathbf{x}~\mathcal{N}(\mathbf{0}, ..