[JAX] 학습한 모델 저장 및 로드
·
연구 Research/인공지능 Artificial Intelligent
기존에 자주 사용되면 파이토치나 텐서플로에서도 그렇듯이 학습한 모델을 저장하는 것은 필수이다. 나중에 다시 결과를 출력해야하거나 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_..