Fonction magique numpy ?

Bonjour tout le monde,

une fois n'est pas coutume, je viens ici poser une question technique. En réalité ce n'est ni urgent ni critique car j'ai déjà une solution suffisament performante pour mes besoins court et moyen terme mais je m'interoge vraiment si on ne peut pas faire mieux.

Mon but est de calculer quelques valeurs dont voici la formulation mathématique (tiré d'un article) :

$$ S_k^0 = \sum_{n=1}^{N} {\gamma_n(k)} $$

$$ S_k^1 = \sum_{n=1}^{N} {\gamma_n(k) * x_n} $$

$$ S_k^2 = \sum_{n=1}^{N} {\gamma_n(k) * x_n^2} $$

Sachant que $x_n \in \mathbb{R}^D$, on a ainsi $S_k^0 \in \mathbb{R}$, $S_k^1 \in \mathbb{R}^D$ et $S_k^2 \in \mathbb{R}^D$. Le $x_n^2$ dans $S_k^2$ est en faite "la mise au carrée de chaque composante".

Bien donc on voit en réalité qu'il y a 3 dimensions dans cette histoire : N, K et D car ce qui m'interesse c'est de calculé ça pour tous les K.

En terme matriciel, avec des boucles, ça peut donner ça :

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21

import numpy as np

def calc_1(gamma, x):
    # gamma.shape == (N,K)
    # x.shape == (N, D)
    N, K = gamma.shape
    _, D = x.shape

    S0 = np.zeros((K,))
    S1 = np.zeros((K, D))
    S2 = np.zeros((K, D))

    for n in range(N):
        for k in range(K):
            gnk = gamma[n, k]
            S0[k] += gnk                    # gamma[n, k]
            for d in range(D):
                tsn = gnk * x[n, d]
                S1[k, d] += tsn             # gamma[n, k] * x[n, d]
                S2[k, d] += tsn * x[n, d]   # gamma[n, k] * x[n, d] * x[n, d]
    return S0, S1, S2

Cette version est évidement ultra lente. J'ai donc fais une version purrement vectoriel en utilisant le broadcasting pour éviter les boucles. Ça donne cela :

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

def calc_2(gamma, x):
    # gamma.shape == (N,K)
    # x.shape == (N, D)

    S0 = np.sum(gamma, axis=0)
    ts = gamma[:, :, np.newaxis] * x[:, np.newaxis, :]
    S1 = np.sum(ts, axis=0)
    S2 = np.sum(ts * x[:, np.newaxis, :], axis=0)

    return S0, S1, S2

Là, en théorie, il y a pas grand chose qui coince… sauf que je vie dans le monde réel. Ce que je n'ai pas encore détaillé c'est les valeurs des dimensions : $K=1024$, $D=96$ et $N=20000$. (autant K et D sont fixe, autant N varie. 20000 est une valeur typique).

Les plus observateurs ont peut être trouvé ce qui me dérange avec calc_2 : la variable intermédiaire ts qui est de dimension $K \times D \times N$. Avec des float sur 64 bits et mes valeurs typique cela nous fais un petit tableau de $T = 1024*96*20000*8/(1024^3) = 14.64$ Go. Et là ça commence à être problématique.

Donc une solution intermédiaire et facile est de mixer les deux :

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17

def calc_3(gamma, x):
    # gamma.shape == (N,K)
    # x.shape == (N, D)

    N, K = gamma.shape
    _, D = x.shape

    S0 = np.sum(gamma, axis=0)
    S1 = np.zeros((K, D))
    S2 = np.zeros((K, D))

    for n in range(N):
        tsn = gamma[n, :, np.newaxis] * x[n, np.newaxis, :]
        S1 += tsn
        S2 += tsn * x[n, np.newaxis, :]

    return S0, S1, S2

On en arrive donc a ma question. Je me demande si il n'y aurait pas une fonction numpy qui ferait tout seul ce genre de calcul, simple, pour m'éviter de déplier une boucle sans exploser ma mémoire. Si vous avez une fonction magique, n'hésitez pas !

Pour vous donner une idée, voici des temps de traitements de ces fonctions (avec toutes les dimensions divisés par 4 pour garder le même ordre de grandeur) :

On se rend donc compte que la version sans boucle n'est pas la plus rapide pour autant. La création et l'accès mémoire doit y être pour beaucoup.

Et le code complet de test, avec les fonctions si vous voulez essayer chez vous :

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70

import numpy as np


def calc_1(gamma, x):
    # gamma.shape == (N,K)
    # x.shape == (N, D)
    N, K = gamma.shape
    _, D = x.shape

    S0 = np.zeros((K,))
    S1 = np.zeros((K, D))
    S2 = np.zeros((K, D))

    for n in range(N):
        for k in range(K):
            gnk = gamma[n, k]
            S0[k] += gnk                    # gamma[n, k]
            for d in range(D):
                tsn = gnk * x[n, d]
                S1[k, d] += tsn             # gamma[n, k] * x[n, d]
                S2[k, d] += tsn * x[n, d]   # gamma[n, k] * x[n, d] * x[n, d]
    return S0, S1, S2


def calc_2(gamma, x):
    # gamma.shape == (N,K)
    # x.shape == (N, D)

    S0 = np.sum(gamma, axis=0)
    ts = gamma[:, :, np.newaxis] * x[:, np.newaxis, :]
    S1 = np.sum(ts, axis=0)
    S2 = np.sum(ts * x[:, np.newaxis, :], axis=0)

    return S0, S1, S2


def calc_3(gamma, x):
    # gamma.shape == (N,K)
    # x.shape == (N, D)

    N, K = gamma.shape
    _, D = x.shape

    S0 = np.sum(gamma, axis=0)
    S1 = np.zeros((K, D))
    S2 = np.zeros((K, D))

    for n in range(N):
        tsn = gamma[n, :, np.newaxis] * x[n, np.newaxis, :]
        S1 += tsn
        S2 += tsn * x[n, np.newaxis, :]

    return S0, S1, S2

if __name__ == "__main__":

    import timeit
    
    # calc_1 n'est calculé que sur 3 itération parce que beaucoup trop lente !
    for i, n in enumerate((3, 100, 100), start=1):
        setup = ";".join(("from __main__ import calc_{}".format(i),
                          "import numpy as np",
                          "np.random.seed(42)",
                          "K = 256",
                          "D = 24",
                          "N = 5000",
                          "gamma = np.random.rand(N, K)",
                          "x = np.random.rand(N, D)"))

        print(i, ": ", timeit.timeit("calc_{}(gamma, x)".format(i), setup=setup, number=n))

Ok bon je vais me répondre moi même.

En réalité la fonction magique que je cherche c'est … le produit matriciel Mes fonctions sont équivalentes à :

1
2
3
4
5
6
7
8
9

def calc_4(gamma, x):
    # gamma.shape == (N,K)
    # x.shape == (N, D)

    S0 = np.sum(gamma, axis=0)
    S1 = np.dot(gamma.transpose(), x)
    S2 = np.dot(gamma.transpose(), np.square(x))

    return S0, S1, S2

tout simplement.

Par contre ça reste toujours plus lent que de déplier la boucle, probablement grace à la mutualisation du calcul intermédiaire entre S1 et S2 :

Donc a moins que quelqu'un a une autre idée, je vais juste conclure de ce sujet que reprendre un problème quelques semaines après et l'exposer par écrit permet généralement de le résoudre seul.

Bon et comme j'aime bien parler et me répondre seul, la conclusion de la conclusion est "de toute façon, Cython écrase tout"

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29

import numpy as np
cimport numpy as np
DTYPE = np.float64
ctypedef np.float64_t DTYPE_t

cimport cython

@cython.boundscheck(False)
@cython.wraparound(False)
def calc_5(np.ndarray[DTYPE_t, ndim=2, mode="c"] gamma, np.ndarray[DTYPE_t, ndim=2, mode="c"] x):
    # Some size constants
    cdef unsigned int N = gamma.shape[0]
    cdef unsigned int K = gamma.shape[1]
    cdef unsigned int D = x.shape[1]

    cdef np.ndarray[DTYPE_t, ndim=1] S0 = np.sum(gamma, axis=0)
    cdef np.ndarray[DTYPE_t, ndim=2] S1 = np.zeros((K, D), dtype=DTYPE)
    cdef np.ndarray[DTYPE_t, ndim=2] S2 = np.zeros((K, D), dtype=DTYPE)
    cdef unsigned int n, k, d
    cdef DTYPE_t tsn, gnk

    for n in range(N):
        for k in range(K):
            gnk = gamma[n, k]
            for d in range(D):
                tsn = gnk * x[n, d]
                S1[k, d] += tsn
                S2[k, d] += tsn * x[n, d]
    return S0, S1, S2

Bon je crois que je vais arrêter là…

Du coup, je n'ai pas vraiment résisté, et voici la version julia:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82

function row_major(gamma, x)
    const N, K = size(gamma)
    const D = size(x, 2)


    S0 = sum(gamma, 1)
    S1 = zeros(Float64, K, D)
    S2 = zeros(Float64, K, D)
    gnk = 0.0
    tsn = 0.0

    @inbounds begin
        for n=1:N
            for k=1:K
                gnk = gamma[n, k]
                for d=1:D
                    tsn = gnk * x[n, d]
                    S1[k, d] += tsn
                    S2[k, d] += tsn * x[n, d]
                end
            end
        end
    end
    return (S0, S1, S2)
end

function column_major(gamma, x)
    const K, N = size(gamma)
    const D = size(x, 1)

    S0 = sum(gamma, 2)
    S1 = zeros(Float64, D, K)
    S2 = zeros(Float64, D, K)
    gnk = 0.0
    tsn = 0.0

    @inbounds begin
        for n=1:N
            for k=1:K
                gnk = gamma[k, n]
                for d=1:D
                    tsn = gnk * x[d, n]
                    S1[d, k] += tsn
                    S2[d, k] += tsn * x[d, n]
                end
            end
        end
    end
    return (S0, S1, S2)
end

function matrix_mult(gamma, x)
    S0 = sum(gamma, 2)
    S1 = gamma' * x
    S2 = gamma' * x.^2
    return (S0, S1, S2)
end

K=1024
D=96
N=20000
small = randn(2,2)

x = randn(N, D)
gamma = randn(N, K)
row_major(small, small)
println("Row-major order")
@time row_major(gamma, x)

x = randn(D, N)
gamma = randn(K, N)

column_major(small, small)
println("Column-major order")
@time column_major(gamma, x)

x = randn(N, D)
gamma = randn(N, K)

matrix_mult(small, small)
println("Matrix multiply")
@time matrix_mult(gamma, x)

Résultats:

1
2
3
4
5
6

Row-major order
elapsed time: 29.752669212 seconds (1 MB allocated)
Column-major order
elapsed time: 5.593091756 seconds (1 MB allocated)
Matrix multiply
elapsed time: 0.315586425 seconds (16 MB allocated, 0.35% gc time in 1 pauses with 0 full sweep)

Donc dans les temps de Numpy bien utilisé, mais pas de Cython. Ou alors il me manque des bouts d'optimisation.

Et la version FORTRAN:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35

program Numpy_bench
    implicit none

    integer, parameter :: N=20000, K=1024, D=96
    double precision, dimension(N, K) :: gamma
    double precision, dimension(N, D) :: x

    double precision, dimension(K) :: S0
    double precision, dimension(K, D) :: S1, S2

    double precision :: t0, t

    S0 = 0
    S1 = 0
    S2 = 0

    call random_number(gamma)
    call random_number(x)

    call cpu_time(t0)

    S0 = sum(gamma, 1)
    S1 = matmul(transpose(gamma), x)
    S2 = matmul(transpose(gamma), square(x))

    call cpu_time(t)
    write(*, *) "Matrix multiply: ", t - t0

contains
    elemental function square(x)
        double precision :: square
        double precision, intent(in) :: x
        square = x*x
    end function
end program

Avec gfortran, et -O3, le calcul met 3s. Comme quoi, BLAS et LAPACK, c'est le bien !

3s ?

J'avoue ne rien connaître à Fortran, mais je dois reconnaître que je suis un peu surpris.
Ce langage a longtemps été une référence dans le calcul scientifique (j'ignore si c'est toujours le cas aujourd'hui, qui l'a remplacé ? C++, Python, un autre ?), c'est curieux d'avoir de si maigres performances (comparé aux autres benchs).

3s avec gfortran, et sans utiliser BLAS. Numpy et Julia utilisent BLAS et gagnent donc pas mal de temps avec ça.

Je testerai ce que donne ifort lundi, là je n'y ai plus accès. Par contre, je suis plus étonné par les perfs de Cython, qui envoient du paté. Il faudrait que je teste une version avec des boucles de fortran.

(j'ignore si c'est toujours le cas aujourd'hui, qui l'a remplacé ? C++, Python, un autre ?)

Il est toujours là, fier et droit ^^. Mais il partage en effet son territoire avec C++, Python, et par selon les domaines R, Mathematica, …

@kje,

Il est clair que la clé pour optimiser ton code est cython, tu peux encore je pense mieux faire avec les boucles, voir la doc…

En tout cas numpy + cython, difficile de faire mieux, les performances sont sans doute très proche du C comme d'habitude.

3s avec gfortran, et sans utiliser BLAS. Numpy et Julia utilisent BLAS et gagnent donc pas mal de temps avec ça.

Je testerai ce que donne ifort lundi, là je n'y ai plus accès. Par contre, je suis plus étonné par les perfs de Cython, qui envoient du paté. Il faudrait que je teste une version avec des boucles de fortran.

ifort et -O3, le temps d'exécution de la version matricielle en fortran tombe à 0.8s. Sur la même machine, gfortran est a 3.4s. Mais ifort est connu pour être particulièrement agressif dans ses optimisations.

Ha marrant ces retours, au moins je n'aurais pas été seul dessus…

@fred1599 :

tu peux encore je pense mieux faire avec les boucles, voir la doc…

Je ne pense pas. J'avais bien creuser la doc tu t'en doute. Et le liens que tu m'a donné le dit clairement :

If i is declared as a cdef integer type, it will optimise this into a pure C loop.

Mes variables de boucles sont bien déclaré en unsigned int et j'avais vérifié le code généré, j'ai bien des boucles C classiques. Pas besoin de la forme "pyrex". C'est de toute façon le plus important. Le gain apporté par le typage de la variable de boucle est très important.