使用Cython提高Python函数的性能

Question

Aim

I'm trying to speed up my Python program with Cython. The code I'm writing is an attempt at the Forward algorithm, used to recursively and efficiently calculate probabilities of long sequences in a Hidden Markov Model (HMM). This problem is usually referred to as the Evaluation Problem.

Python Code

In a file called hmm.py

import numpy
import pandas

class HMM():
    '''    
    args:
        O:
            observation sequence. A list of 'H's or 'T's

        X:
            state sequence. 'S','M' or 'L's

        A:
            transition matrix, N by N

        B:
            Emission matrix, M by N

        M:
            Number of possibilities in emission matrix

        pi:
            initial transition matrix

        N:
            Number of states

        T:
            length of the observation sequence

        Q:
            possible hidden states (Xs)

        V:
            possible observations (Os)

    '''
    def __init__(self,A,B,pi,O,X):
        self.A=A
        self.N=self.A.shape[0]
        self.B=B
        self.M=self.B.shape[1]
        self.pi=pi
        self.O=O
        self.T=len(O)
        self.Q=list(self.A.index)
        self.V=list(self.B.keys())
        self.X=X


    def evaluate(self):
        '''
        Solve the evaluation problem for HMMs 
        by implementing the forward algorithm
        '''
        c0=0
        ct=numpy.zeros(self.T)
        alpha= numpy.zeros((self.T,self.N))

        ## compute alpha[0]
        for i in range(self.N):
            pi0=self.pi[self.Q[i]]
            bi0=self.B.loc[self.Q[i]][self.O[0]]
            alpha[0,i]=pi0*bi0
            c0+=alpha[0,i]
            ct[0]=alpha[0,i]
        ct[0]=1/ct[0]#[0]
        alpha=alpha*ct[0]
        for t in range(1,self.T):
            for i in range(self.N):
                for j in range(self.N):
                    aji= self.A[self.Q[j]][self.Q[i]]
                    alpha[t,j]= alpha[t-1,j]*aji
                ct[t]=ct[t]+alpha[t,i]
            ct[t]=1/ct[t]
            alpha=alpha*ct[t]
        return (alpha,ct)  

This code can be called with:

if __name__=='__main__':
    A=[[0.7,0.3],[0.4,0.6]]
    A= numpy.matrix(A)
    A=pandas.DataFrame(A,columns=['H','C'],index=['H','C'])
    '''
    three types of emmission, small s, medium m and large l
    '''
    B=[[0.1,0.4,0.5],[0.7,0.2,0.1]]
    B=numpy.matrix(B)
    B=pandas.DataFrame(B,columns=['S','M','L'],index=['H','C'])
    '''
    initial probabilities for state, H and C
    '''
    pi=[0.6,0.4]
    pi=numpy.matrix(pi)
    pi=pandas.DataFrame(pi,columns=['H','C'])
    O=(0,1,0,2)
    O=('S','M','S','L')
    X=('H','H','C','C')
    H=HMM(A,B,pi,O,X)
    print H.evaluate()

When using %timeit I get this output with pure python

Compilation with Cython

I then placed the evaluate function (rather than the whole class) in a new file hmm.pyx extension:

import numpy 
cimport numpy

cpdef evaluate_compiled(A,B,pi,O,X):
    '''
    Solve the evaluation problem for HMMs 
    by implementing the forward algorithm
    '''
    T=len(O)
    N=len(list(set(X)))
    Q=list(set(X))
    V=list(set(O))

    c0=0
    ct=numpy.zeros(T)
    alpha= numpy.zeros((T,N))

    ## compute alpha[0]
    for i in range(N):
        pi0=pi[Q[i]]
        bi0=B.loc[Q[i]][O[0]]
        alpha[0,i]=pi0*bi0
        c0+=alpha[0,i]
        ct[0]=alpha[0,i]
    ct[0]=1/ct[0]#[0]
    alpha=alpha*ct[0]
    for t in range(1,T):
        for i in range(N):
            for j in range(N):
                aji= A[Q[j]][Q[i]]
                alpha[t,j]= alpha[t-1,j]*aji
            ct[t]=ct[t]+alpha[t,i]
        ct[t]=1/ct[t]
        alpha=alpha*ct[t]
    return (alpha,ct)   

In setup.py :

try:
    from setuptools import setup
    from setuptools import Extension
except ImportError:
    from distutils.core import setup
    from distutils.extension import Extension

from Cython.Distutils import build_ext
import numpy

setup(cmdclass={'build_ext':build_ext},
      ext_modules=[Extension('hmm',
                             sources=['HMMCluster/hmm.pyx'], #File is in a directory called HMMCluster
                                     include_dirs=[numpy.get_include()])]   )

Now after compilation, I can use:

from hmm import evaluate_compiled

And under the __main__ block above I can use, inplace of evaluate :

print evaluate_compiled(A,B,pi,O,X)

and with %timeit :

As you can see, without changing the code I have a ~3 fold improvement in speed. However, all the doc's I have read suggest the lack of speed in Python is due to dynamically inferring variable types. Therefore, in principle, I can declare variable types and speed things up further.

Cython with type declaration

Now, the last function on this post is the same algorithm again but with type declaration

cpdef evaluate_compiled_with_type_declaration(A,B,pi,O,X):
    cdef int t,i,j
    cdef int T  = len(O)
    cdef int N  = len(list(set(X)))
    cdef list Q = list(set(X))
    cdef list V = list(set(O))
    cdef float c0 = 0
#    cdef numpy.ndarray ct = numpy.zeros(T,dtype=double) ## this caused compilation to fail
    ct=numpy.zeros(T)
    alpha= numpy.zeros((T,N))
    for i in range(N):
        pi0=pi[Q[i]]
        bi0=B.loc[Q[i]][O[0]]
        alpha[0,i]=pi0*bi0
        c0+=alpha[0,i]
        ct[0]=alpha[0,i]
    ct[0]=1/ct[0]#[0]
    alpha=alpha*ct[0]
    for t in range(1,T):
        for i in range(N):
            for j in range(N):
                aji= A[Q[j]][Q[i]]
                alpha[t,j]= alpha[t-1,j]*aji
            ct[t]=ct[t]+alpha[t,i]
        ct[t]=1/ct[t]
        alpha=alpha*ct[t]
    return (alpha,ct)   

After compilation and %timeit i get:

As you can see, the type declarations haven't made any further improvements to the code performance. My question is: Can any further improvements be made to the speed of this code and if so, how do I do it?

Edit Following suggestions in the comments I added the additional type declarations:

cdef float  pi0,bi0
cdef numpy.ndarray[numpy.float64_t, ndim=1] ct
cdef numpy.ndarray[numpy.float64_t, ndim=2] aij,alpha

and got this from %timeit :

So again, still no real speedups even with declaring types.

Answer 1

The "modern" way to use NumPy arrays in Cython is "Typed Memoryviews" http://docs.cython.org/en/latest/src/userguide/memoryviews.html

The corresponding arguments must be declared as:

cpdef evaluate_compiled_with_type_declaration(double[:,:] A, double[:,:] B, double[:] pi, int[:] O, int[:] X):

(just guessing for types and shapes).

They must be indexed directly as A[i,j] and not A[i][j]

You can declare your result arrays and fill them in one go.

    cdef double[:] ct=numpy.zeros(T)
    cdef double[:,:] alpha= numpy.zeros((T,N))

For optimization, see the compiler directives http://cython.readthedocs.io/en/latest/src/reference/compilation.html?highlight=cdivision#compiler-directives (specifically cdivision and boundscheck )

You should not convert your input data via list(set(...)) as the set loses ordering.

To check whether your code is "compiled", use cython -a on the file as suggested by Warren Weckesser