huggingface transformers pipeline gpu 지정해서 모델 load하는 방법

예시)

 

2, 3번 gpu 에만 모델 업로드하고 싶다 그러면 import os 를 최상단에 해주고 visible devices를 아래와 같이 정의하면 된다.(정확히는 최상단에 해주기보다 gpu가 사용되는 라이브러리 import 전에 해줄 것!) 

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "2, 3"

import transformers
...

 

꼭 저렇게 하드코딩 할 필요는 없어 보이고 python실행할때 붙여줘도 될 거 같다.

CUDA_VISIBLE_DEVICES=2,3 python ~

 

device_map  은 그냥 "auto" 로 놔두면 된다.

 

더 좋은 방법이 있을 거 같은데 일단 이렇게 해서 됐다.