使用gpu与cpu运行tensorflow 时间比较

使用gpu与cpu运行tensorflow 时间比较

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
热身时间
cpu_time=timeit.timeit(cpu_run,number=10)
gpu_time=timeit.timeit(gpu_run,number=10)
print('warmup time',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)

笔记本gtx2060运行结果如下：