用tensorflow,pytorch框架使用GPU,指定GPU问题

本文详细介绍了如何在PyTorch和TensorFlow中指定和管理GPU资源,包括单个GPU的指定使用方法、多GPU的并行处理策略以及如何限制GPU内存使用量,适合深度学习开发者参考。

摘要生成于 C知道 ,由 DeepSeek-R1 满血版支持, 前往体验 >

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()  # 返回当前设备索引

pytorch DataParallel 多GPU使用 

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) # 列出所有的本地机器设备

 

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值