[PyTorch Lightning] Trainer 클래스에서 GPU 설정하는 방법

 

 

 

PyTorch Lightning에서 Trainer 클래스를 사용할 때, GPU 설정하는 방법에 대해 알아보겠습니다.

 

1. accelerator 설정

GPU를 사용하려면, accelerator 파라미터를 "gpu"로 설정하면 됩니다.

trainer = pl.Trainer(accelerator="gpu")

 

 

2. devices 설정

devices 파라미터를 이용해서, 학습에 사용할 GPU 개수를 설정하거나, 인덱스로 특정 GPU를 지정할 수도 있습니다.

 

(1) GPU 개수 설정

trainer = pl.Trainer(devices=2, accelerator="gpu")

2개 GPU를 사용해서 학습할 때는, 위와 같이 설정하면 됩니다.

 

(2) 인덱스로 특정 GPU 지정

trainer = pl.Trainer(devices=[1, 3], accelerator="gpu")

특정 GPU를 학습에 사용하고 싶을 때는, 위와 같이 GPU 인덱스를 리스트 형태로 설정해 주면 됩니다.

 

(3) 사용 가능한 모든 GPU 사용

trainer = pl.Trainer(devices=-1, accelerator="gpu")

사용 가능한 모든 GPU를 사용하고 싶을 때는, 위와 같이 devices 파라미터를 -1로 설정해 주면 됩니다.

 

(4) 자동 설정

trainer = pl.Trainer(devices="auto", accelerator="gpu")

자동으로 설정해 주는 방법도 있습니다. 현재 사용 가능한 GPU 정보를 잘 모를 때 사용하면 유용할 것 같습니다.

default 값은 "auto"입니다.

 

 

제가 작성한 내용은 아래 PyTorch Lightning 공식 홈페이지 내용을 참고해서 작성하였습니다.

링크에 들어가 보시면, 더 자세한 내용과 다른 옵션에 대한 설명도 확인하실 수 있습니다.

https://lightning.ai/docs/pytorch/stable/common/trainer.html