pytorch:
单个GPU指定使用:
1.在代码上指定(建议这种指定方式)
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "2" # 指定第二块gpu
2.使用set_device指定
import torch
torch.cuda.set_device(0) # 不能设定两个,比如0,2
查看pytorch所使用的gpu
torch.cuda.is_available() # cuda是否可用
torch.cuda.device_count() # 返回gpu数量
torch.cuda.get_device_name(0) # 返回gpu名字,设备索引默认从0开始
torch.cuda.current_device() # 返回当前设备索引
tensorflow:
1.指定gpu
import os
os.environ["CUDA_VISIBLE_DEVICES"]=‘1’
2.设置定量的GPU使用量
gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.7)
sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options)) #占用现存不超过总现存的0.7
3.按需设置最小显存
gpu_options = tf.GPUOptions(allow_growth=True)
sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))
查看tensorflow所使用的gpu
import tensorflow as tf
tf.test.gpu_device_name() # 显示当前使用gpu名字
tf.test.is_gpu_available() # true or false
from tensorflow.python.client import device_lib
local_device_protos = device_lib.list_local_devices()
print(local_device_protos) # 列出所有的本地机器设备