tensorflow2.0的cpu与gpu运行时间对比

文章目录

  • 前言
  • 一、导入环境
  • 二、定义函数
  • 三、测试


前言

这里运用一个自定义大小的矩阵数据计算,来测试gpu与cpu运算时间的对比。


以下为实现方法

一、导入环境

示例:pandas 是基于NumPy 的一种工具,该工具是为了解决数据分析任务而创建的。

#设置显卡内存使用率,根据使用率占用
import os
os.environ["TF_FORCE_GPU_ALLOW_GROWTH"] = "true"
import tensorflow as tf
print('Tensorflow:{}'.format(tf.__version__))
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline
import timeit as timeit

二、定义函数

#创建在cpu环境上运算的两个矩阵
n=100
with tf.device('/cpu:0'):
    cpu_a = tf.random.normal([1,n])
    cpu_b = tf.random.normal([n,1])
    print(cpu_a.device,cpu_b.device)
    
#创建在gpu环境上运算的两个矩阵
with tf.device('/gpu:0'):
    gpu_a = tf.random.normal([1,n])
    gpu_b = tf.random.normal([n,1])
    print(gpu_a.device,gpu_b.device)
    
def cpu_run(): #cpu运算时间
    with tf.device('/cpu:0'):
        c = tf.matmul(cpu_a,cpu_b)
    return c
def gpu_run(): #cpu运算时间
    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:',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)

你可能感兴趣的:(人工智能,python,tensorflow)