CuPy:加速Numpy中的矩阵运算

前言

在计算机视觉任务中,经常要使用python中的numpy模块对图像进行预处理,但是numpy是通过调用cpu进行运算,在运算速度上相比gpu还是逊色不少。恰巧,笔者在工作中遇到了强化学习中的environment在返回state和reward时占用了大量的训练时间这样的问题。将代码中的numpy替换为cupy后,网络的训练速度提升了近40倍,直接起飞~

Numpy vs CuPy

import numpy as np
import cupy as cp
import time

array_np = np.ones((512,512,3))
array_cp = cp.ones((512,512,3))

# numpy
np_start = time.time()
for _ in range(1000):
    c = np.power(array_np, 2)
np_end = time.time()

# cupy
cp_start = time.time()
for _ in range(1000):
    c = cp.power(array_cp, 2)
cp_end = time.time()

print("Calculate time with Numpy: {}".format(np_end - np_start))
print("Calculate time with CuPy: {}".format(cp_end - cp_start))

Calculate time with Numpy: 11.282422065734863
Calculate time with CuPy: 0.05162358283996582

由运行结果可以看到,加速效果真的是非常的amazing呀!

猜你喜欢

转载自blog.csdn.net/zxdd2018/article/details/124724428
今日推荐