python代码设置gpu
时间: 2025-06-20 20:44:12 浏览: 16
### 如何在Python代码中配置和使用GPU
#### TensorFlow中的GPU配置方法
为了使TensorFlow能够识别并使用GPU,在编写代码之前需确认已安装了支持GPU的TensorFlow版本[^2]。可以通过如下方式验证当前环境是否成功启用了GPU:
```python
import tensorflow as tf
# 检查tensorflow能否访问GPU
print("Num GPUs Available: ", len(tf.config.experimental.list_physical_devices('GPU')))
```
如果上述命令返回大于0的数量,则表示环境中至少有一块可用的GPU设备。
对于希望指定特定GPU执行任务的情况,可以采用下面这段代码实现对单个或多个GPU的选择性分配:
```python
gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
try:
# 设置仅限于第一个GPU可见
tf.config.set_visible_devices(gpus[0], 'GPU')
# 或者限制显存增长以防止占用全部内存
for gpu in gpus:
tf.config.experimental.set_memory_growth(gpu, True)
except RuntimeError as e:
print(e)
```
以上操作允许开发者灵活控制程序运行时所依赖的具体硬件资源[^1]。
#### PyTorch中的GPU配置方法
同样地,在PyTorch框架下也可以轻松完成类似的设定过程。首先检测是否有可用的CUDA设备:
```python
import torch
# 判断是否存在可用的cuda设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f'Using device {device}')
```
当需要将模型迁移到GPU上训练时,只需简单调用`.to(device)`即可完成迁移工作;同理,数据张量也应被转移到相同的设备之上以便参与运算:
```python
model.to(device) # 将模型移动到选定设备
input_tensor = input_tensor.to(device) # 同样处理输入张量
output = model(input_tensor) # 执行前向传播
```
通过这种方式可以在不改变原有逻辑结构的前提下充分利用现代图形处理器的强大性能提升计算效率。
阅读全文
相关推荐

















