[딥러닝] 파이토치 개발환경 구축 2. CUDA로 GPU 사용하기

r4bb1t.tistory.com/59

[딥러닝] 파이토치 개발환경 구축

개발환경을 구축했으면 이제 프로그램을 GPU에서 돌려볼 수 있습니다. 설치가 모두 제대로 되었는지 확인해봅시다.

 

import torch

print(torch.cuda.is_available())

 

제대로 설치가 됐다면 True가 뜹니다.

 

import torch

CUDA_AVAILABLE = torch.cuda.is_available()
DEVICE = torch.device("cuda" if CUDA_AVAILABLE else "cpu")

'''모델 선언'''

model = ModelExample().to(DEVICE)

 

이렇게 삼항연산자를 사용해서 DEVICE를 정의해주면 CUDA 사용이 가능한 환경에서는 DEVICE는 GPU로, 아닐 경우 CPU를 의미합니다. 모델을 선언할 때 .to(DEVICE)를 달아줍니다.

 

받아오는 데이터들은

 

data = data.to(DEVICE)

 

등으로 한번 더 처리를 해줘야합니다.