[GN] Penzai - 신경망 구축, 편집, 시각화를 위한 JAX 도구 키트 (feat. Google Deepmind)

Penzai - 신경망 구축, 편집, 시각화를 위한 JAX 도구 키트

소개

  • 구글 DeepMind가 발표한, 모델을 읽기 쉽고 기능적인 Pytree 데이터 구조로 작성하기 위한 JAX 라이브러리
  • 모델이 학습된 후에도 쉽게 작업할 수 있도록 하는 데 중점을 둠
    • 모델 구성 요소를 리버스 엔지니어링 또는 제거 및 내부 활성화 검사 및 조사, 모델 수술, 아키텍처 디버깅 등과 관련된 연구에 적합
    • 모델 시각화, 수정 및 분석용 도구를 제공해서, 그냥 모델을 만들고 훈련하는 용도로도 사용 가능
  • 함께 설계되었지만 각각 독립적으로 사용 가능한 모듈형 도구 모음으로 구성

포함된 도구들

penzai.nn (pz.nn)

  • Flax, Haiku, Keras, Equinox 등의 다른 신경망 라이브러리에 대한 대안으로, 선언적 컴비네이터 기반 신경망 라이브러리
  • 모델의 forward pass 전체 구조를 모델 pytree에 노출하여, pretty printing으로 모델의 모든 작업을 볼 수 있고 jax.tree_util로 새로운 런타임 로직을 주입할 수 있음
  • 모델은 내부적으로 호출 가능한 pytree에 불과함

penzai.treescope (pz.ts)

  • 일반 IPython/Colab 렌더러를 대체하는 드롭인 방식으로 작동하는 강력한 대화형 Python pretty-printer
  • Penzai 모델과 기타 깊이 중첩된 JAX pytree를 이해하는 데 도움이 되도록 설계되었으며, 임의 차원의 NDArray 시각화를 기본적으로 지원함

penzai.core.selectors (pz.select)

  • JAX의 .at[...].set(...) 구문을 임의의 유형 기반 pytree 순회로 일반화하는 pytree 용 다기능 도구
  • Penzai 모델 및 기타 데이터 구조의 복잡한 재작성 또는 온더플라이 패치를 쉽게 수행할 수 있음

penzai.core.named_axes (pz.nx)

  • 일반 JAX 함수를 명명된 축에 대해 벡터화하고, 새로운 배열 API를 배우지 않고도 명명된 프로그래밍 스타일과 위치 프로그래밍 스타일 간에 원활하게 전환할 수 있는 경량 명명된 축 시스템

penzai.data_effects (pz.de)

  • pytree 순회를 기반으로 구축되고 모델 작성이나 사용을 방해하지 않으면서 제어권을 갖는, 부수 인수, 난수 및 상태 변수에 대한 옵트인 시스템

원문

Penzai GitHub 저장소

Penzai 튜토리얼

https://penzai.readthedocs.io/

출처 / GeekNews


:information_source: 알려드립니다

이 글은 국내외 IT 소식들을 공유하는 GeekNews의 운영자이신 xguru님께 허락을 받아 GeekNews에 게제된 AI 관련된 소식을 공유한 것입니다.

출처의 GeekNews 링크를 방문하시면 이 글과 관련한 추가적인 의견들을 보시거나 공유하실 수 있습니다! :wink:

:gift: 아래:arrow_lower_right:쪽에 좋아요:heart:를 눌러주시면 새로운 소식을 정리하고 공유하는데 힘이 됩니다~ :star_struck: