Created
October 12, 2021 06:24
-
-
Save ugo-nama-kun/e3105dd66a0e0c35a7aa9e9949b4a6b3 to your computer and use it in GitHub Desktop.
tensorflow で GPU のリソースを制限するやつ
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import tensorflow as tf # chacked @ tensorflow==2.6.0 | |
available_gpus = tf.config.experimental.list_physical_devices('GPU') | |
print("Num GPUs Available: ", len(available_gpus)) | |
if available_gpus: | |
try: | |
tf.config.experimental.set_visible_devices(available_gpus[gpu_id], "GPU") | |
logical_gpus = tf.config.experimental.list_logical_devices('GPU') | |
print(len(available_gpus), "Physical GPUs,", len(logical_gpus), "Logical GPU") | |
except RuntimeError as e: | |
print(e) |
PyTorchだと以下のような感じ
https://deideeplearning.com/2020/04/03/pytorch-gpu/
https://pystyle.info/pytorch-how-to-specify-the-device-for-calculation/
https://note.nkmk.me/python-pytorch-cuda-is-available-device-count/
def get_device(gpu_id=-1):
if gpu_id >= 0 and torch.cuda.is_available():
return torch.device("cuda", gpu_id)
else:
return torch.device("cpu")
device = get_device()
print(device) # cpu
device = get_device(gpu_id=0)
print(device) # cuda:0
import torch
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
A = torch.randn(3, 5)
A.to(device)
print(torch.cuda.get_device_name(0))
# GeForce GTX 1080 Ti
print(torch.cuda.get_device_name(torch.device('cuda:0')))
# GeForce GTX 1080 Ti
print(torch.cuda.get_device_name('cuda:0'))
# GeForce GTX 1080 Ti
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
import torch
print(torch.cuda.is_available())
# False
print(torch.cuda.device_count())
# 0
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
from https://www.tensorflow.org/guide/gpu?hl=ja