如何使用numba加速多维Longsumexp和softmax

Question

我想使用多项式 logit 计算 softmax/概率，同时使用

longsumexp

以避免溢出。使用 numba 带来 2-3 倍的加速。我在这里可以做得更好吗？另外，当我使用

fastmath=True

时，它似乎没有带来任何加速，那么我是否以错误的方式编写了

numba

循环？

import numba
import numpy as np
def get_p_4d(a, lamda):
    m = a * lamda[:, None][:,None].transpose(0,3,1,2)
    c = np.max(m, axis=2)[:,None].transpose(0,2,1,3)
    aa = np.exp(m - c)
    logsumexp = c + np.log(aa.sum(axis=2)[:,None].transpose(0,2,1,3))
    p = np.exp(m - logsumexp)
    return p

@numba.njit()
def get_p_4d_nb(a, lamda, num_code, num_draw, num_action):
    p = np.empty((num_code, num_draw, num_action, 3))
    a = a.transpose(0, 1, 3, 2)
    for i in range(num_code):
        for j in range(num_draw):
            this_lamda = lamda[i,j]
            for k in range(num_action):
                p[i, j, k, 0] = a[i, j, k, 0] * this_lamda
                p[i, j, k, 1] = a[i, j, k, 1] * this_lamda
                p[i, j, k, 2] = a[i, j, k, 2] * this_lamda

                c = p[i,j,k,0]
                c = max(c, p[i,j,k,1])
                c = max(c, p[i,j,k,2])

                logsumexp = np.log(
                    np.exp(p[i, j, k, 0] - c) + np.exp(p[i, j, k, 1] - c) + np.exp(p[i, j, k, 2] - c)) + c

                p[i, j, k, 0] = np.exp(p[i, j, k, 0] - logsumexp)
                p[i, j, k, 1] = np.exp(p[i, j, k, 1] - logsumexp)
                p[i, j, k, 2] = np.exp(p[i, j, k, 2] - logsumexp)

    return p.transpose(0, 1, 3, 2)

a=np.ones((112,1000,3,3))
lamda = np.random.uniform(0., 1., size=112*1000).reshape(112,1000)
get_p_4d(a, lamda)
get_p_4d_nb(a, lamda, 112, 1000, 3)

Answer 1

您可以尝试并行化任务（我还使用切片减少了代码

0:3

）：

@numba.njit(parallel=True)
def get_p_4d_nb_parallel(a, lamda, num_code, num_draw, num_action):
    p = np.empty((num_code, num_draw, num_action, 3), dtype="float32")
    a = a.transpose(0, 1, 3, 2)
    for i in numba.prange(num_code):
        for j in range(num_draw):
            this_lamda = lamda[i, j]
            for k in range(num_action):
                p[i, j, k, 0:3] = a[i, j, k, 0:3] * this_lamda
                c = np.max(p[i, j, k, 0:3])
                logsumexp = np.log(np.exp(p[i, j, k, 0:3] - c).sum()) + c
                p[i, j, k, 0:3] = np.exp(p[i, j, k, 0:3] - logsumexp)
    return p.transpose(0, 1, 3, 2)

---

基准：

from timeit import timeit

import numba
import numpy as np


def get_p_4d(a, lamda):
    m = a * lamda[:, None][:, None].transpose(0, 3, 1, 2)
    c = np.max(m, axis=2)[:, None].transpose(0, 2, 1, 3)
    aa = np.exp(m - c)
    logsumexp = c + np.log(aa.sum(axis=2)[:, None].transpose(0, 2, 1, 3))
    p = np.exp(m - logsumexp)
    return p


@numba.njit
def get_p_4d_nb(a, lamda, num_code, num_draw, num_action):
    p = np.empty((num_code, num_draw, num_action, 3))
    a = a.transpose(0, 1, 3, 2)
    for i in range(num_code):
        for j in range(num_draw):
            this_lamda = lamda[i, j]
            for k in range(num_action):
                p[i, j, k, 0] = a[i, j, k, 0] * this_lamda
                p[i, j, k, 1] = a[i, j, k, 1] * this_lamda
                p[i, j, k, 2] = a[i, j, k, 2] * this_lamda

                c = p[i, j, k, 0]
                c = max(c, p[i, j, k, 1])
                c = max(c, p[i, j, k, 2])

                logsumexp = (
                    np.log(
                        np.exp(p[i, j, k, 0] - c)
                        + np.exp(p[i, j, k, 1] - c)
                        + np.exp(p[i, j, k, 2] - c)
                    )
                    + c
                )

                p[i, j, k, 0] = np.exp(p[i, j, k, 0] - logsumexp)
                p[i, j, k, 1] = np.exp(p[i, j, k, 1] - logsumexp)
                p[i, j, k, 2] = np.exp(p[i, j, k, 2] - logsumexp)

    return p.transpose(0, 1, 3, 2)


@numba.njit(parallel=True)
def get_p_4d_nb_parallel(a, lamda, num_code, num_draw, num_action):
    p = np.empty((num_code, num_draw, num_action, 3), dtype="float32")
    a = a.transpose(0, 1, 3, 2)
    for i in numba.prange(num_code):
        for j in range(num_draw):
            this_lamda = lamda[i, j]
            for k in range(num_action):
                p[i, j, k, 0:3] = a[i, j, k, 0:3] * this_lamda
                c = np.max(p[i, j, k, 0:3])
                logsumexp = np.log(np.exp(p[i, j, k, 0:3] - c).sum()) + c
                p[i, j, k, 0:3] = np.exp(p[i, j, k, 0:3] - logsumexp)
    return p.transpose(0, 1, 3, 2)


a = np.ones((112, 1000, 3, 3))
lamda = np.random.uniform(0.0, 1.0, size=112 * 1000).reshape(112, 1000)

x = get_p_4d(a, lamda)
y = get_p_4d_nb(a, lamda, 112, 1000, 3)
z = get_p_4d_nb_parallel(a, lamda, 112, 1000, 3)

assert np.allclose(x, y)
assert np.allclose(x, z)


t1 = timeit("get_p_4d(a, lamda)", number=1, globals=globals())
t2 = timeit("get_p_4d_nb(a, lamda, 112, 1000, 3)", number=1, globals=globals())
t3 = timeit("get_p_4d_nb_parallel(a, lamda, 112, 1000, 3)", number=1, globals=globals())

print(t1, t2, t3, sep="\n")

在我的机器上打印（AMD 5700x）：

0.033313663996523246
0.01112839400593657
0.0017921309990924783