tensorflow で GPU のリソースを制限するやつ

tf_gpu_resource.py

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)

from https://www.tensorflow.org/guide/gpu?hl=ja

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