segunda-feira, 7 de dezembro de 2015

Acelerando Python usando Numba

Implementações em Python puro são frequentemente muito lentas e qualquer tentativa de acelerar estes programas impacta diretamente a produtividade de um cientista que uma Python. Falamos anteriormente de Processamento Paralelo com multiprocessing e Cython como maneiras de acelerar experimentos. Ambas soluções, entretanto, exigem modificações não triviais (e potencialmente extensivas) no código, aumentando sua complexidade.

Neste texto irei apresentar uma alternativa chamada Numba, um compilador Just in Time(JIT) que compila um subconjunto da linguagem Python em código de máquina eficiente. Numba é muito mais limitado que Cython e não é capaz de traduzir todas (a maioria das?) construções Python de maneira eficiente. As funcionalidades suportadas, porém, podem ser aceleradas significativamente sem grande esforço.

Primeiramente, um compilador JIT traduz, durante a execução, código escrito em uma linguagem de alto nível como Python em instruções nativas da CPU em que está sendo executado. Otimizações são aplicadas somente nas partes mais lentas do código. A vantagem disto é que não é necessário compilar nada antes da execução do programa, porém as primeira execuções de uma função "jitted" são mais lentas.

Uma das grandes vantagens de se usar Numba é que o código compilado é Python puro. Não é necessário adicionar nenhum tipo de anotação ou código extra. Tudo é feito automaticamente pelo Numba e o código gerado é bastante rápido em alguns casos específicos. A maneira mais fácil de instalar este pacote é utilizando o Anaconda Python. Um simples conda install numba instala todo o necessário sem nenhuma complicação, mesmo em Windows.

Vamos usar como exemplo nosso já cansado gradiente morfológico. Também manteremos a versão Cython que criamos em um post anterior.

Versão Python

In [1]:
import numpy as np
import scipy as sp
import scipy.ndimage
import time
import matplotlib.pyplot as plt

def morpho_gradient_py(img, out, p):
    pi, pj = p
    minv = 255
    maxv = 0
    
    for i in range(-1, 2):
        for j in range(-1, 2):
            if img[pi+i, pj+j] < minv:
                minv = img[pi+i, pj+j]
            if img[pi+i, pj+j] > maxv:
                maxv = img[pi+i, pj+j]
    out[pi, pj] = maxv - minv

def versao_python(img):
    out = np.zeros(img.shape, img.dtype)
    for i in range(1, img.shape[0]-1):
        for j in range(1, img.shape[1]-1):
            morpho_gradient_py(img, out, (i, j))
    return out

Versão Cython

In [2]:
%load_ext cython
In [3]:
%%cython
import numpy as np
import scipy as sp
import scipy.ndimage
import time
import matplotlib.pyplot as plt
import cython

@cython.boundscheck(False)
@cython.nonecheck(False)
cdef morpho_gradient_cython(long[:,:] img, long[:,:] out, p):
    cdef int i, j, pi, pj, minv, maxv
    
    pi, pj = p
    minv = 255
    maxv = 0
    
    for i in range(-1, 2):
        for j in range(-1, 2):
            if img[pi+i, pj+j] < minv:
                minv = img[pi+i, pj+j]
            if img[pi+i, pj+j] > maxv:
                maxv = img[pi+i, pj+j]
    out[pi, pj] = maxv - minv

@cython.boundscheck(False)
@cython.nonecheck(False)
cpdef versao_cython(long[:,:] img):
    cdef int i, j, w, h
    cdef long[:,:] out
    h = img.shape[0]; w = img.shape[1]
    out = np.zeros((h, w), np.int)
    for i in range(1, h-1):
        for j in range(1, w-1):
            morpho_gradient_cython(img, out, (i, j))
    return out

Versão Numba

In [4]:
import numba

@numba.jit
def morpho_gradient_numba(img, out, p):
    pi, pj = p
    minv = 255
    maxv = 0
    
    for i in range(-1, 2):
        for j in range(-1, 2):
            if img[pi+i, pj+j] < minv:
                minv = img[pi+i, pj+j]
            if img[pi+i, pj+j] > maxv:
                maxv = img[pi+i, pj+j]
    out[pi, pj] = maxv - minv

@numba.jit
def versao_numba(img):
    out = np.zeros(img.shape, img.dtype)
    for i in range(1, img.shape[0]-1):
        for j in range(1, img.shape[1]-1):
            morpho_gradient_numba(img, out, (i, j))
    return out

A versão Numba é basicamente a versão Python com um decorador adicionado às funções. Diferentemente de Cython, não adicionamos nenhum tipo manualmente nem precisamos compilar nosso código explicitamente. Vamos agora aos resultados numéricos.

In [5]:
import scipy

img = scipy.misc.lena()

print('Versão Python')
%timeit versao_python(img)
print('Versão Cython')
%timeit versao_cython(img)
print('Versão Numba')
%timeit versao_numba(img)
Versão Python
1 loops, best of 3: 3.23 s per loop
Versão Cython
10 loops, best of 3: 39.7 ms per loop
Versão Numba
The slowest run took 25.12 times longer than the fastest. This could mean that an intermediate result is being cached 
1 loops, best of 3: 26.6 ms per loop

Ambas as versões aceleradas obtiveram performance muito superior ao programa em Python puro. Surpreendentemente, o código em Numba obteve um desempenho superior ao código em Cython. Isto ocorre, principalmente, pois usamos somente funcionalidades suportadas no modo nativo do Numba, chamado de nopython. Isto inclui, basicamente, manipulação de arrays com Numpy (funções suportadas), chamadas de função com decorador @jit e controle de fluxo padrão de Python (if, for i in range(..), while). Existe suporte limitado para listas e é melhor evitar funções recursivas. Uma lista completa das funcionalidades suportadas pode ser vista na documentação do projeto. Funções nopython podem ser compiladas para código nativo bastante rápido e o próprio processo de compilação não só é invisível para o programador com também é bem rápido, diferentemente de Cython.

Porém, o modo nopython só é ativado se a função só utilizar as funcionalidades permitidas. Se uma única operação não for suportada o Numba ativa o modo object e trata todas as variáveis como objetos Python. Ou seja, apesar de algumas otimizações ainda serem possíveis, o programa será muito mais lento. Não é possível, como em Cython, misturar código que interage com o interpretador e código compilado. Vejamos abaixo um exemplo do mesmo código acima compilado no modo object.

In [6]:
@numba.jit(forceobj=True)
def morpho_gradient_numba_ruim(img, out, p):
    pi, pj = p
    minv = 255
    maxv = 0
    
    for i in range(-1, 2):
        for j in range(-1, 2):
            if img[pi+i, pj+j] < minv:
                minv = img[pi+i, pj+j]
            if img[pi+i, pj+j] > maxv:
                maxv = img[pi+i, pj+j]
    out[pi, pj] = maxv - minv

@numba.jit(forceobj=True)
def versao_numba_ruim(img):
    out = np.zeros(img.shape, img.dtype)
    for i in range(1, img.shape[0]-1):
        for j in range(1, img.shape[1]-1):
            morpho_gradient_numba_ruim(img, out, (i, j))
    return out


print('Versão Numba Object')
%timeit versao_numba_ruim(img)
Versão Numba Object
1 loops, best of 3: 18.3 s per loop

Como podemos ver, o resultado pode ser desastroso. Apesar da aparente limitação, Numba é muito útil para o que se propõe a fazer: acelerar código em Python que contém diversos loops e acessa arrays do Numpy. Uma boa maneira de checar se um código poderá ser acelerado é usar o decorador @jit(nopython=True). Esta opção força a compilação em modo nopython e levanta uma exceção se o código não puder ser compilado deste modo.

Mesmo que a função inteira não possa ser compilada, se houver algum loop que só possui operações nopython o Numba pode executar uma otimização chamada loop lifting. Vejamos o exemplo abaixo.

In [7]:
@numba.jit
def morpho_gradient_numba_lift(img, out, p):
    pi, pj = p
    minv = 255
    maxv = 0
    for i in range(-1, 2):
        for j in range(-1, 2):
            if img[pi+i, pj+j] < minv:
                minv = img[pi+i, pj+j]
            if img[pi+i, pj+j] > maxv:
                maxv = img[pi+i, pj+j]
    out[pi, pj] = maxv - minv

@numba.jit
def versao_numba_lift(img):
    out = np.zeros(img.shape, np.uint8)
    d = {}

    for i in range(1, img.shape[0]-1):
        for j in range(1, img.shape[1]-1):
            morpho_gradient_numba_lift(img, out, (i, j))

    d[5] = 0

    return out


print('Versão Numba Object')
%timeit versao_numba_lift(img)
Versão Numba Object
The slowest run took 19.17 times longer than the fastest. This could mean that an intermediate result is being cached 
1 loops, best of 3: 25.3 ms per loop

A função versao_numba_lift não pode ser compilada em modo nopython, mas os dois for podem. Logo, ele extrai esta parte da função, aplica loop lifting e termina com um código tão rápido quanto o anterior. Porém, isto não funcionaria bem se fizéssemos o mesmo na função morpho_gradient_numba_lift. Apesar do loop da outra função ser realmente acelerado, uma funcionalidade importante do Numba é que a chamada de funções em modo nopython é muito rápida para outras funções nopython, mas lenta para funções Python (ref). Como em Cython, é possível ganhar muita velocidade se todas as nossas funções forem compiladas no modo acelerado.


Como pudemos ver, Numba pode ser uma excelente alternativa para acelerar código escrito em Python. Em um próximo post apresentarei as funções (e limitações) de Numba para rodar código em GPUs NVidia usando CUDA.

Um comentário:

  1. por favor perdoe-me a ignorância, mas com o numba da pra gerar um executável tipo .exe como as linguagens mais conhecidas. dipo delphi vb etc. obrigado

    ResponderExcluir