본인의 컴퓨터나 노트북에 NVIDIA GPU가 있고 CUDA가 이미 설치되어있는 것을 전제로 한다.또한 가상환경을 conda에서 만들어서 그 안에서 최신 버전의 JAX를 설치하는 과정이다. (일반적으로 JAX를 pip으로 설치하면 무조건 최신버전을 설치하도록 된다.) 설치한 JAX 버전은 현재 기준으로 가장 최신인 0.9.0이다. 이전 버전을 설치하고 싶다면 다음 글이 더 유용하다. 나는 주로 Linux x86_64, Windows WSL2 환경을 사용한다. 다음 커맨드를 실행하면 JAX GPU 버전이 설치가 되는데 여기서 가정하는 것은 CUDA가 적절하게 설치되어있다는 것이다.pip install --upgrade "jax[cuda12]" 최근에는 JAX installation 문서가 잘 업데이트..

jupyer lab/notebook을 쓰다가 kernel이 터지는 경우가 있는데 다음과 같은 경우들이 있다. https://github.com/microsoft/vscode-jupyter/wiki/Kernel-crashes Kernel crashesVS Code Jupyter extension. Contribute to microsoft/vscode-jupyter development by creating an account on GitHub.github.com  내가 겪은 대표적인 원인들은 다음과 같다.1. 라이브러리 설치 실패 - 호환 안됨 등의 문제2. gpu에 업로드한 데이터가 너무 많아서 문제  1번의 경우에는 재설치하고 버전 호환을 신경써서 설치해야한다. 아나콘다 등의 버전 관리 시스템을 쓰면..

Quasi Newton method를 이용해서 JAX에서 최적화를 시키고 싶었는데 JAX 자체가 Deep learning에 포커스가 있고 대부분 딥러닝이 Gradient descent method로 최적화를 하다보니 라이브러리를 찾게 되었다. 대부분 jaxopt라는 라이브러리를 추천했기 때문에 이걸로 수행해보았다. jaxopt에는 jaxopt.ScipyMinimize와 jaxopt.LBFGS가 있는데 다른 분들의 시도를 보니 ScipyMinimize가 더 성능이 괜찮은 것 같다. ScipyMinimize는 scipy에 있는 최적화를 사용한 것이고 LBFGS는 직접 만든 것 같은데 line search 방법 등이 다르다고 한다. import jax import jax.numpy as np from jax ..

inject_hyperparams라는 함수로 optax의 optimizer를 묶어서 사용하면 hyperparams를 관찰할 수 있다. # Wrap the optimizer to inject the hyperparameters optimizer = optax.inject_hyperparams(optax.adamw)(learning_rate=schedule) def fit(params: optax.Params, optimizer: optax.GradientTransformation) -> optax.Params: opt_state = optimizer.init(params) # Since we injected hyperparams, we can access them directly here print(f'A..

JAX에서 값을 쓰다가 어려운 부분이 있으면 대체로 numpy에 있는 함수들과 비슷한 것이 많아서 찾기 쉽다. 1. NaN 값 찾기 import jax.numpy as jnp a = jnp.array([jnp.nan,1,0,jnp.nan]) x = jnp.isnan(a) print(x) >> [ True False False True] 2. +-Inf 값 찾기 import jax.numpy as jnp a = jnp.array([jnp.inf,1,0,-jnp.inf]) x = jnp.isinf(a) print(x) >> [ True False False True] 1,2번 항목을 보면 boolean array으로 나오기 때문에 위 코드에서 'a'라는 array에 대해 indexing하면 NaN값을 추출할..

기존에 자주 사용되면 파이토치나 텐서플로에서도 그렇듯이 학습한 모델을 저장하는 것은 필수이다. 나중에 다시 결과를 출력해야하거나 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를 쓰다가 너무 많은 양의 데이터를 쓰다보니 메모리 부족(OOM: Out of memory) 현상을 겪었다. 근본적인 해결책은 달리 없다. 데이터가 너무 많아서 생기는 문제이니 데이터 양을 줄이던가 아니면 병렬 컴퓨팅을 하는 방법이 있다. 병렬 컴퓨팅을 간단한 코드에서는 실행해보았는데 큰 네트워크에서는 해본 적이 없다. 일단 임시방편으로는 다음과 같다. 메모리 부족이 쉽게 발생할 수 있는 이유는 JAX에서 처음에 import를 할 때 대부분의 메모리를 미리 할당해놓기 때문이다. 따라서 이 preallocation을 막거나 줄이면 도움이 된다. 1. Preallocation 중단 XLA_PYTHON_CLIENT_PREALLOCATE=false # 구체적으로는 다음과 같이 구현한다. import os ..

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

가장 simple하게 신경망을 구성하는 방법에 대해서 저장해놓은 글이다.차츰 업데이트 할 예정 1. 기본 학습 코드import jaximport jax.numpy as jnpfrom jax import grad, jit, vmapfrom jax import random# Define a simple neural network modeldef init_params(layer_sizes, key): params = [] for i in range(1, len(layer_sizes)): key, subkey = random.split(key) w = random.normal(subkey, (layer_sizes[i-1], layer_sizes[i])) b = j..