[JAX] device 확인, default device 설정

2023. 4. 13. 16:38·연구 Research/인공지능 Artificial Intelligent

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를 import하기 전에

import os
os.environ['CUDA_VISIBLE_DEVICES'] = "2"

이렇게 인식할 수 있는 devices를 지정한다.

 

import jax
nmp = jax.numpy.ones(4)
print(nmp.device())
print(jax.devices())
>> gpu:0
[GpuDevice(id=0, process_index=0)]

그 다음에 배열을 만들어서 확인을 했는데 device를 조회해보면 id=0 밖에 인식이 되지 않는다.

gpu:0이지만 실제로는 내가 지정한 세 번째 gpu이다.

이를 확인하는 방법은 nvidia-smi를 터미널에 입력하는 것이다.

 

 

다음과 같이 JAX는 처음에 import할 때 90퍼센트 가까이 preallocation하기 때문에 위와 같이 세 번째 gpu의 메모리가 차지된 것을 확인할 수 있다.

저작자표시 비영리 변경금지 (새창열림)

'연구 Research > 인공지능 Artificial Intelligent' 카테고리의 다른 글

[JAX] 메모리 부족 문제 해결  (1) 2023.04.26
[인공지능] Ubuntu 18.04에서 CUDA, CuDNN 설치  (1) 2023.04.16
[JAX] JAX와 Torch, CUDA, cudnn 버전 맞추기  (0) 2023.04.12
[JAX] 학습 중 NaN 값이 나올 때 찾는 방법  (0) 2023.03.28
[JAX] Gaussian process 파라미터에 따른 결과 visualization  (0) 2023.03.24
'연구 Research/인공지능 Artificial Intelligent' 카테고리의 다른 글
  • [JAX] 메모리 부족 문제 해결
  • [인공지능] Ubuntu 18.04에서 CUDA, CuDNN 설치
  • [JAX] JAX와 Torch, CUDA, cudnn 버전 맞추기
  • [JAX] 학습 중 NaN 값이 나올 때 찾는 방법
보통의공대생
보통의공대생
수학,프로그래밍,기계항공우주 등 공부하는 기록들을 남깁니다.
  • 보통의공대생
    뛰는 놈 위에 나는 공대생
    보통의공대생
  • 전체
    오늘
    어제
    • 분류 전체보기 (468)
      • 공지 (1)
      • 영어 공부 English Study (40)
        • 텝스 TEPS (7)
        • 글 Article (21)
        • 영상 Video (10)
      • 연구 Research (99)
        • 최적화 Optimization (3)
        • 데이터과학 Data Science (7)
        • 인공지능 Artificial Intelligent (40)
        • 제어 Control (45)
      • 프로그래밍 Programming (103)
        • 매트랩 MATLAB (25)
        • 파이썬 Python (33)
        • 줄리아 Julia (2)
        • C++ (3)
        • 리눅스 우분투 Ubuntu (6)
      • 항공우주 Aeronautical engineeri.. (21)
        • 항법 Navigation (0)
        • 유도 Guidance (0)
      • 기계공학 Mechanical engineering (13)
        • 열역학 Thermodynamics (0)
        • 고체역학 Statics & Solid mechan.. (10)
        • 동역학 Dynamics (1)
        • 유체역학 Fluid Dynamics (0)
      • 수학 Mathematics (34)
        • 선형대수학 Linear Algebra (18)
        • 미분방정식 Differential Equation (3)
        • 확률및통계 Probability & Sta.. (2)
        • 미적분학 Calculus (1)
        • 복소해석학 Complex Analysis (5)
        • 실해석학 Real Analysis (0)
      • 수치해석 Numerical Analysis (27)
      • 확률 및 랜덤프로세스 Random process (2)
      • 추론 & 추정 이론 Estimation (3)
      • 기타 (26)
        • 설계 프로젝트 System Design (8)
        • 논문작성 Writing (55)
        • 세미나 Seminar (2)
        • 생산성 Productivity (3)
      • 실험 Experiment (1)
      • 유학 생활 Daily (8)
  • 블로그 메뉴

    • 홈
    • 태그
    • 방명록
  • 링크

  • 공지사항

  • 인기 글

  • 태그

    Dear abby
    에러기록
    옵시디언
    텝스
    고체역학
    Python
    우분투
    Statics
    ChatGPT
    Linear algebra
    생산성
    논문작성
    teps
    IEEE
    JAX
    pytorch
    LaTeX
    서버
    obsidian
    Julia
    텝스공부
    Numerical Analysis
    논문작성법
    Zotero
    MATLAB
    matplotlib
    인공지능
    WOX
    수치해석
    딥러닝
  • 최근 댓글

  • 최근 글

  • hELLO· Designed By정상우.v4.10.3
보통의공대생
[JAX] device 확인, default device 설정
상단으로

티스토리툴바