import tensorflow as tf
import timeit # 性能分析模块
# 查看版本
print(tf.__version__)
# CUDA是否可用
print(tf.test.is_built_with_cuda())
# GPU是否可用
print(tf.test.is_gpu_available())
# 创建在cpu环境上运算的两个矩阵
# tf.device:指定模型运行的具体设备
with tf.device('/cpu:0'):
# 用于从“服从指定正态分布的序列”中随机取出指定个数的值
# shape:输出张量的形状,必选
# mean:正态分布的均值,默认为0
# stddev:正态分布的标准差,默认为1.0
# dtype:输出的类型,默认为tf.float32
# seed:随机数种子,是一个整数,当设置之后,每次生成的随机数都一样
# name:操作的名称
cpu_a = tf.random.normal(shape=[10000, 1000])
cpu_b = tf.random.normal([1000, 2000])
print(cpu_a.device, cpu_b.device)
# 创建在gpu环境上运算的两个矩阵
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)
# cpu矩阵相乘
def cpu_run():
with tf.device('/cpu:0'):
# 将矩阵a乘以矩阵b,生成a*b
c = tf.matmul(cpu_a, cpu_b)
return c
# gpu矩阵相乘
def gpu_run():
with tf.device('/gpu:0'):
c = tf.matmul(gpu_a, gpu_b)
return c
# 正式计算前需要热身,避免将初始化时间结算在内
cpu_time = timeit.timeit(cpu_run, number=10)
gpu_time = timeit.timeit(gpu_run, number=10)
print('warm up:', 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)
正确运行结果
2.4.1
True
True
/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
warm up: 0.6571273 1.2346933
run time: 0.6092833999999998 0.001225699999999108
可见GPU运算速度要比CPU快得多