Skip to content

Latest commit

 

History

History
67 lines (46 loc) · 1.94 KB

File metadata and controls

67 lines (46 loc) · 1.94 KB

GPU计算

到目前为止,我们一直在使用CPU计算。对复杂的神经网络和大规模的数据来说,使用CPU来计算可能不够高效。在本节中,我们将介绍如何使用单块NVIDIA GPU来计算。首先,需要确保已经安装好了至少一块NVIDIA GPU。然后,下载CUDA并按照提示设置好相应的路径

注意:需要tensorflow-gpu(注:在tensorflow2以上的版本不需要直接安装gpu版本,因为不区分gpu和cpu,有gpu可用时会自动调用gpu)

import tensorflow as tf
import numpy as np
print(tf.__version__)

gpus = tf.config.experimental.list_physical_devices(device_type='GPU')
cpus = tf.config.experimental.list_physical_devices(device_type='CPU')
print("可用的GPU:",gpus,"\n可用的CPU:", cpus)
2.0.0
可用的GPU: [PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')] 
可用的CPU: [PhysicalDevice(name='/physical_device:CPU:0', device_type='CPU')]

check available device

from tensorflow.python.client import device_lib
print(device_lib.list_local_devices())
[name: "/device:CPU:0"
device_type: "CPU"
memory_limit: 268435456
locality {
}
incarnation: 15592483132577835191
, name: "/device:GPU:0"
device_type: "GPU"
memory_limit: 3063309926
locality {
  bus_id: 1
  links {
  }
}
incarnation: 3859243074925251015
physical_device_desc: "device: 0, name: GeForce GTX 1650, pci bus id: 0000:01:00.0, compute capability: 7.5"
]

specify device

使用tf.device()来指定特定设备(GPU/CPU)

with tf.device('GPU:0'):
    a = tf.constant([1,2,3],dtype=tf.float32)
    b = tf.random.uniform((3,))
    print(tf.exp(a + b) * 2)
    
tf.Tensor([12.172885 19.682476 53.18001 ], shape=(3,), dtype=float32)

注:本节与原书有很大不同,原书传送门