PyTorch를 사용하여 총 GPU 메모리와 사용 가능한 메모리량을 얻는 방법
PyTorch를 사용하여 총 GPU 메모리와 사용 가능한 메모리량을 얻는 방법
필요한 라이브러리
- Python
- PyTorch
- CUDA (선택 사항)
코드
import torch
# 전체 GPU 메모리 용량 확인
total_mem = torch.cuda.get_device_properties(0).total_memory
# 사용 가능한 GPU 메모리 용량 확인
free_mem = torch.cuda.memory_allocated(0)
# 사용 중인 GPU 메모리 용량 확인
used_mem = total_mem - free_mem
# 메모리 정보 출력
print(f"총 GPU 메모리: {total_mem / 1024**3:.3f} GB")
print(f"사용 가능한 GPU 메모리: {free_mem / 1024**3:.3f} GB")
print(f"사용 중인 GPU 메모리: {used_mem / 1024**3:.3f} GB")
코드 설명
torch.cuda.get_device_properties(0)
: 첫 번째 GPU 장치의 속성을 가져옵니다.total_memory
: 전체 GPU 메모리 용량 (바이트)torch.cuda.memory_allocated(0)
: 첫 번째 GPU 장치에서 할당된 메모리 용량 (바이트)used_mem
: 사용 중인 GPU 메모리 용량 (바이트)
참고
- CUDA를 설치하지 않은 경우 CPU 메모리만 사용할 수 있습니다.
- 여러 GPU 장치가 있는 경우
torch.cuda.get_device_properties(i)
를 사용하여 다른 장치의 메모리 정보를 확인할 수 있습니다. torch.cuda.memory_allocated(i)
는 특정 GPU 장치에서 할당된 메모리만 반환합니다. CPU 메모리 사용량은 포함하지 않습니다.
추가 정보
예제 코드
import torch
# GPU 사용 가능 여부 확인
if torch.cuda.is_available():
# GPU 사용
device = torch.device("cuda")
# GPU 메모리 정보 확인
total_mem = torch.cuda.get_device_properties(device).total_memory
free_mem = torch.cuda.memory_allocated(device)
used_mem = total_mem - free_mem
# 메모리 정보 출력
print(f"총 GPU 메모리: {total_mem / 1024**3:.3f} GB")
print(f"사용 가능한 GPU 메모리: {free_mem / 1024**3:.3f} GB")
print(f"사용 중인 GPU 메모리: {used_mem / 1024**3:.3f} GB")
else:
# CPU 사용
device = torch.device("cpu")
# CPU 메모리 정보 확인
total_mem = psutil.cpu_info().total_memory
free_mem = psutil.virtual_memory().available
used_mem = total_mem - free_mem
# 메모리 정보 출력
print(f"총 CPU 메모리: {total_mem / 1024**3:.3f} GB")
print(f"사용 가능한 CPU 메모리: {free_mem / 1024**3:.3f} GB")
print(f"사용 중인 CPU 메모리: {used_mem / 1024**3:.3f} GB")
설명
- 이 코드는 GPU 사용 가능 여부를 확인하고 사용 가능한 경우 GPU 메모리 정보를 출력합니다.
- GPU를 사용할 수 없는 경우 CPU 메모리 정보를 출력합니다.
psutil
라이브러리는 CPU 및 메모리 사용량과 같은 시스템 정보를 가져오는 데 사용됩니다.
참고
psutil
라이브러리는 설치해야 사용할 수 있습니다.
총 GPU 메모리와 사용 가능한 메모리량을 얻는 다른 방법
nvidia-smi 명령어 사용
nvidia-smi
nvidia-smi
명령어는 GPU 정보를 출력하는 명령어입니다.- 출력 결과에서
Total Memory
와Free Memory
항목을 확인하면 총 GPU 메모리와 사용 가능한 메모리량을 알 수 있습니다.
시스템 정보 도구 사용
PyTorch 라이브러리 사용
import torch
# 사용 가능한 GPU 장치 목록 확인
devices = torch.cuda.device_count()
# 각 GPU 장치의 메모리 정보 확인
for i in range(devices):
total_mem = torch.cuda.get_device_properties(i).total_memory
free_mem = torch.cuda.memory_allocated(i)
used_mem = total_mem - free_mem
# 메모리 정보 출력
print(f"** GPU 장치 {i} **")
print(f"총 GPU 메모리: {total_mem / 1024**3:.3f} GB")
print(f"사용 가능한 GPU 메모리: {free_mem / 1024**3:.3f} GB")
print(f"사용 중인 GPU 메모리: {used_mem / 1024**3:.3f} GB")
- 이 코드는 사용 가능한 모든 GPU 장치의 메모리 정보를 출력합니다.
참고
nvidia-smi
명령어는 CUDA가 설치되어 있어야 사용할 수 있습니다.- 시스템 정보 도구는 운영 체제에 따라 다릅니다.
결론
본 문서에서는 Python, PyTorch, GPU를 사용하여 총 GPU 메모리와 사용 가능한 메모리량을 얻는 방법을 설명했습니다. 또한,
nvidia-smi
명령어 사용- 시스템 정보 도구 사용
- PyTorch 라이브러리 사용
python pytorch gpu