机器学习加速利器jax，让numpy加速30倍

梦里梦外; 2023-07-15 09:14 10阅读 0赞

**jax.numpy是CPU、GPU和TPU上的numpy，具有出色的自动差异化功能，可用于高性能机器学习研究**。

我今天就来试一试到底多快。我在同一台bu带gpu的机器上进行试验

首先我们得安装jax

pip install jax jaxlib

先试一下原生的numpy

import numpy as np  
    import time
    x = np.random.random([5000, 5000]).astype(np.float32)
    try:
        st=time.time()
        y=np.matmul(x, x)
    except Exception:
        print("erro")
    print(time.time()-st)
    print(y)

运行结果：

[root@node opt]# python np.py
    4.424036026000977
    [[1236.3004 1240.3048 1211.4501 ... 1225.7804 1237.1368 1235.1566]
     [1235.5778 1246.7327 1208.7142 ... 1238.117  1232.439  1226.5779]
     [1235.0111 1244.4628 1211.5264 ... 1238.5541 1246.9045 1244.6909]
     ...
     [1229.7677 1238.8345 1210.4467 ... 1219.8604 1234.0862 1220.1482]
     [1231.9464 1251.9636 1212.1384 ... 1235.8513 1236.8677 1240.5355]
     [1254.0636 1265.74   1241.6528 ... 1245.015  1259.153  1247.0613]]

再来试一下jax带的numpy

import jax.numpy as np
    from jax import random
    import time
    x = random.uniform(random.PRNGKey(0), [5000, 5000])
    st=time.time()
    try:
       y=np.matmul(x, x)
    except Exception:
        print("erro")
    print(time.time()-st)
    print(y)

结果：

[root@node opt]# python jax_np.py
    /opt/AN/lib/python3.7/site-packages/jax/lib/xla_bridge.py:122: UserWarning: No GPU/TPU found, falling back to CPU.
      warnings.warn('No GPU/TPU found, falling back to CPU.')
    0.013895750045776367
    [[1261.0647 1244.5797 1237.2269 ... 1264.7208 1246.0367 1260.5391]
     [1256.1    1239.737  1237.5562 ... 1257.1333 1243.5856 1243.5979]
     [1261.2687 1239.5006 1250.6697 ... 1259.8387 1250.6825 1248.5712]
     ...
     [1265.9805 1230.9077 1244.4961 ... 1264.2374 1241.5995 1244.9274]
     [1262.9971 1253.961  1256.2424 ... 1266.3489 1255.1581 1274.1865]
     [1273.3524 1252.4921 1261.0496 ... 1273.2394 1272.829  1267.7483]]

我们可以看到，没有jax的numpy运行了差不多4.4秒，而带了jax的numpy直接才0.014，速度基本上提升了30倍。也太快了

同样地，jax下面还有jax.scipy等代替原生的scipy

![watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3pob3VfNDM4_size_16_color_FFFFFF_t_70][]

[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3pob3VfNDM4_size_16_color_FFFFFF_t_70]: /images/20230528/33bf0545a3084aa6ab87777d98a4c90f.png