若在虚拟环境中使用 conda 安装，conda 会自动安装相关的 cuda 和 cudnn 依赖。
conda install tensorflow-gpu
下面给出一段 GPU 测试程序：
import time import tensorflow as tf begin = time.time() with tf.device('/gpu:0'): rand_t = tf.random_uniform([50,50],0,10,dtype=tf.float32,seed=0) a = tf.Variable(rand_t) b = tf.Variable(rand_t) c = tf.matmul(a,b) init = tf.global_variables_initializer() sess = tf.Session(config=tf.ConfigProto(log_device_placement=True)) #强制使用GPU sess.run(init) print(sess.run(c)) end = time.time() print(end-begin,'s')
运行结果:
2020-10-11 09:36:56.713905: I tensorflow/core/platform/cpu_feature_guard.cc:142] Your CPU supports instructions that this TensorFlow binary was not compiled to use: AVX AVX2 2020-10-11 09:36:56.742138: I tensorflow/stream_executor/platform/default/dso_loader.cc:42] Successfully opened dynamic library nvcuda.dll 2020-10-11 09:36:57.193243: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1640] Found device 0 with properties: name: GeForce GTX 960M major: 5 minor: 0 memoryClockRate(GHz): 1.176 pciBusID: 0000:01:00.0 2020-10-11 09:36:57.199952: I tensorflow/stream_executor/platform/default/dlopen_checker_stub.cc:25] GPU libraries are statically linked, skip dlopen check. 2020-10-11 09:36:57.204946: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1763] Adding visible gpu devices: 0 2020-10-11 09:36:58.622428: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1181] Device interconnect StreamExecutor with strength 1 edge matrix: 2020-10-11 09:36:58.626880: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1187] 0 2020-10-11 09:36:58.630073: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1200] 0: N 2020-10-11 09:36:58.633910: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1326] Created TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 1382 MB memory) -> physical GPU (device: 0, name: GeForce GTX 960M, pci bus id: 0000:01:00.0, compute capability: 5.0) Device mapping: /job:localhost/replica:0/task:0/device:GPU:0 -> device: 0, name: GeForce GTX 960M, pci bus id: 0000:01:00.0, compute capability: 5.0 2020-10-11 09:36:58.657797: I tensorflow/core/common_runtime/direct_session.cc:296] Device mapping: /job:localhost/replica:0/task:0/device:GPU:0 -> device: 0, name: GeForce GTX 960M, pci bus id: 0000:01:00.0, compute capability: 5.0
import torch flag = torch.cuda.is_available() print(flag)
ngpu= 1 # Decide which device we want to run on device = torch.device("cuda:0"if (torch.cuda.is_available() and ngpu > 0) else"cpu") print(device) print(torch.cuda.get_device_name(0)) print(torch.rand(3,3).cuda())
