tensorflow-gpu-2.1安装成功测试代码

以下代码任选其一

测试代码一：

import tensorflow as tf
import timeit


with tf.device('/cpu:0'):
	cpu_a = tf.random.normal([10000, 1000])
	cpu_b = tf.random.normal([1000, 2000])
	print(cpu_a.device, cpu_b.device)

with tf.device('/gpu:0'):
	gpu_a = tf.random.normal([10000, 1000])
	gpu_b = tf.random.normal([1000, 2000])
	print(gpu_a.device, gpu_b.device)

def cpu_run():
	with tf.device('/cpu:0'):
		c = tf.matmul(cpu_a, cpu_b)
	return c 

def gpu_run():
	with tf.device('/gpu:0'):
		c = tf.matmul(gpu_a, gpu_b)
	return c 


# warm up
cpu_time = timeit.timeit(cpu_run, number=10)
gpu_time = timeit.timeit(gpu_run, number=10)
print('warmup:', cpu_time, gpu_time)


cpu_time = timeit.timeit(cpu_run, number=10)
gpu_time = timeit.timeit(gpu_run, number=10)
print('run time:', cpu_time, gpu_time)

输出类似下面：

 GPU (device: 0, name: GeForce MX250, pci bus id: 0000:02:00.0, compute capability: 6.1)
/job:localhost/replica:0/task:0/device:CPU:0 /job:localhost/replica:0/task:0/device:CPU:0
/job:localhost/replica:0/task:0/device:GPU:0 /job:localhost/replica:0/task:0/device:GPU:0
2020-04-15 00:11:16.634792: I tensorflow/stream_executor/platform/default/dso_loader.cc:44] Successfully opened dynamic library cublas64_10.dll
warmup: 1.4536040000000003 0.4982911999999997
run time: 1.3492356 0.0006344000000000349

测试代码二：

import tensorflow as tf
import os

os.environ['TF_CPP_MIN_LOG_LEVEL']='2'

a = tf.constant(1.)
b = tf.constant(3.)
print(a+b)

print('GPU:', tf.test.is_gpu_available())

输出：

tf.Tensor(4.0, shape=(), dtype=float32)
.
.# 一大堆乱七八糟
.
MX250, pci bus id: 0000:02:00.0, compute capability: 6.1) # GPU不同则不同
GPU: True