sexta-feira, 21 de abril de 2017

Desempenho em Python: Cython vs Numpy vs Numba

Como já abordado em vários outros textos, a utilização de Cython (partes 1, 2, 3 e 4) e Numba (CPU e GPU) pode aumentar muito a velocidade de programas escritos em Python. Para programas que tratam arrays, o Numpy oferece também diversas funções muito rápidas que podem ser usadas para acelerar programas. Neste texto farei uma comparação entre programas em Cython, Numba e Numpy puro para duas tarefas de processamento de imagens.

Usaremos o módulo timeit para comparar os resultados. Ao chamarmos timeit.repeat cada função é executada 10 vezes e estas 10 execuções são repetidas 5 vezes. Isto visa evitar que outros processos atrapalhem uma das execuções e tornem a comparação injusta. Pegamos o menor tempo das execuções como uma expectativa otimista do desempenho de cada função.

1. Contagem de pontos brancos em imagens

Apesar de parece um pouco trivial, este exemplo pode dar uma ideia de qual o tempo necessário para percorrer todos os pontos de uma imagem e executar algum processamento. Uma implementação em Python pode ser vista abaixo.

import numpy as np
import scipy as sp
import scipy.ndimage
import matplotlib.pyplot as plt
import timeit

img = sp.ndimage.imread('mask.gif', mode='L')
plt.imshow(img, cmap='gray')
nexec = 10
repeat = 5

def contar_brancos_puro(img):
    count = 0
    for i in range(img.shape[0]):
        for j in range(img.shape[1]):
            if img[i, j] != 0:
                count += 1
    return count

Abaixo exploramos o uso do Numpy e de duas funções embutidas. A função contar_brancos_np1 cria uma nova matriz binária com True nos pixels diferentes de \(0\) e faz a soma de todos os pixels da matriz. A função contar_brancos_np2 usa a função nonzero para devolver dois arrays contendo as posições não zero da imagem e retorna o tamanho deste array. Ambas as funções possuem apenas uma linha e são bastante legíveis. Também testaremos a função np.count_nonzero, que já faz isto e é embutida no Numpy.

def contar_brancos_np1(img):
    return np.sum(img != 0)

def contar_brancos_np2(img):
    return np.nonzero(img)[0].shape

def contar_brancos_np3(img):
    return np.count_nonzero(img)

O código da função contar_brancos_puro pode ser acelerado por Numba diretamente, pois não usa nenhuma primitiva Python proibida no modo nopython. É importante lembrar que o código em Numba é compilado na primeira execução. Por isso executamos o loop várias vezes e pegamos o menor tempo.

import numba as nb

contar_brancos_nb = nb.jit(contar_brancos_puro, nopython=True)

Abaixo está o código anotado em Cython. Tirando as anotações de tipos, não existem diferenças significativas em relação ao código Python puro.

cimport numpy as np
import numpy as np
cimport cython

@cython.boundscheck(False)
@cython.nonecheck(False)
cpdef int contar_brancos_pyx(unsigned char[:,:] img):
    cdef int count = 0
    for i in range(img.shape[0]):
        for j in range(img.shape[1]):
            if img[i, j] != 0:
                count += 1
    return count

Veja abaixo um resumo dos resultados encontrados.

import pyximport; pyximport.install()
from brancos import contar_brancos_pyx
from collections import OrderedDict

funcs = [('Python', 'contar_brancos_puro'),
         ('Numpy 1', 'contar_brancos_np1'),
         ('Numpy 2', 'contar_brancos_np2'),
         ('Numpy 3', 'contar_brancos_np3'),
         ('Cython', 'contar_brancos_pyx'),
         ('Numba', 'contar_brancos_nb')]

print('''
Função   Min    Max
-------  ------ ------''')
for n, f in funcs:
    values = timeit.repeat('%s(img)'%f, globals=globals(),
number=nexec, repeat=repeat)
    mn = min(values)
    mx = max(values)
    print('%s %0.4f %0.4f'%(n.ljust(8), mn, mx))

print('')
Função   Min    Max
-------  ------ ------
Python   7.9207 12.3363
Numpy 1  0.0090 0.0094
Numpy 2  0.0555 0.0613
Numpy 3  0.0167 0.0173
Cython   0.0047 0.0056
Numba    0.0014 0.2025

O código em Numba obteve o melhor desempenho mínimo, mas também o pior máximo. Isto se dá devido à necessidade de compilar a função em sua primeira utilização. É possível, porém, salvar o resultado de compilações anteriores no disco e evitar este problema. O código em Cython obteve desempenho similar ao usando Numba, mas com menor variância entre maior e menor tempo.

A primeira implementação Numpy é melhor que a segunda, possivelmente pois usar np.nonzero implica alocar memória e preencher a lista de posições. De qualquer maneira, ambos resultados são parecidos. Usar np.count_nonzero resulta em desempenho similar à primeira abordagem.

Finalmente, o código em Python puro é muito pior que todas as outras alternativas. Apesar deste resultado já ser esperado, é um pouco assutador perceber que todas as opções melhoraram o desempenho em ao menos 10 vezes. Em especial, o código em Numba é exatamente o mesmo mas foi 10 vezes mais rápido no pior caso e 1000 vezes no melhor caso.

2. Operações de imagens usando janelas

O segundo problema que analisamos é o deslizamento de uma janela em todos os pontos de uma imagem. Este tipo de código é escrito para criar diversos operadores de imagens simples, como convoluções, dilatações e erosões. Neste exemplo iremos apenas copiar os pontos para um novo array X sem fazer nenhum processamento.

A versão em Python puro escrita abaixo não utiliza nenhuma funcionalidade do Numpy. Isto é um pouco artificial, já que o código fica muito mais longo e menos natural para quem trabalha com imagens. Na seção de Numpy faremos uma utilização incremental de suas funções.

def copia_subjanelas_puro(img, msk, win, X):
    k = 0
    w2 = win.shape[1]//2
    h2 = win.shape[0]//2
    for i in range(h2, img.shape[0]-h2):
        for j in range(w2, img.shape[1]-w2):
            if msk[i, j] != 0:
                k2 = 0
                for l in range(-h2, h2+1):
                    for m in range(-w2, w2+1):
                        X[k, k2] = img[i+l, j+m]
                        k2 += 1
                k += 1
    assert k == X.shape[0]

Abaixo estão algumas versões usando Numpy. A primeira somente faz a cópia dos pixels usando notação de intervalos. A segunda substitui o loop por uma chamada da função np.nonzero. Esperamos que a segunda seja ligeiramente mais rápida devido à menor quantidade de loops em Python.

def copia_subjanelas_np1(img, msk, win, X):
    k = 0
    w2 = win.shape[1]//2
    h2 = win.shape[0]//2
    for i in range(h2, img.shape[0]-h2):
        for j in range(w2, img.shape[1]-w2):
            if msk[i, j] != 0:
                X[k] = img[i-h2:i+h2+1,j-w2:j+w2+1].reshape((1, 49))
                k += 1
    assert k == X.shape[0]

def copia_subjanelas_np2(img, msk, win, X):
    w2 = win.shape[1]//2
    h2 = win.shape[0]//2
    idx_i, idx_j = np.nonzero(msk[h2:-h2,w2:-w2])

    for k in range(idx_i.shape[0]):
        i, j = idx_i[k], idx_j[k]
        X[k] = img[i-h2:i+h2+1,j-w2:j+w2+1].reshape((1, 49))

    assert X.shape[0] == idx_i.shape[0]

Também fazemos duas versões em Numba. A primeira só aplica o jit na versão em Python puro. A segunda usa np.nonzero ao invés dos dois loops encaixados. Um ponto importante que não conseguimos aplicar o jit em nenhuma das funções que fizemos com Numpy. Numba só suporta a operação reshape em arrays contínuos, o que não é o caso quando selecionamos uma subimagem usando a indexação com intervalos.

copia_subjanelas_puro_nb = nb.jit(copia_subjanelas_puro,
nopython=True)

@nb.jit
def copia_subjanelas_nb_np(img, msk, win, X):
    w2 = win.shape[1]//2
    h2 = win.shape[0]//2
    idx_i, idx_j = np.nonzero(msk[h2:-h2,w2:-w2])

    for k in range(idx_i.shape[0]):
        i, j = idx_i[k]+h2, idx_j[k]+w2
        k2 = 0
        for l in range(-h2, h2+1):
            for m in range(-w2, w2+1):
                X[k, k2] = img[i+l, j+m]
                k2 += 1

    assert X.shape[0] == idx_i.shape[0]

A versão em Cython é igual a em Python puro, mas com anotações de tipos.

cimport numpy as np
import numpy as np
cimport cython


@cython.boundscheck(False)
@cython.nonecheck(False)
cpdef void copia_subjanelas_cy(unsigned char[:,:]img, unsigned
char[:,:] msk, unsigned char[:,:] win, unsigned char[:,:] X):
    cdef int h2 = win.shape[0]//2
    cdef int w2 = win.shape[1]//2
    cdef int i, j, l, m, k, k2

    k = 0
    for i in range(h2, img.shape[0]-h2):
        for j in range(w2, img.shape[1]-w2):
            if msk[i, j] != 0:
                k2 = 0
                for l in range(-h2, h2+1):
                    for m in range(-w2, w2+1):
                        X[k, k2] = img[i+l, j+m]
                        k2 += 1
                k+=1
    assert k == X.shape[0]
from subjanelas import copia_subjanelas_cy


msk = img
img = sp.ndimage.imread('input.tif', mode='L')
win = np.ones((7, 7), np.uint8)
nmsk = contar_brancos_nb(msk)
X = np.zeros((nmsk, 49), np.uint8)

print('Número de subjanelas copiadas', X.shape[0])

funcs = [('Python', 'copia_subjanelas_puro'),
         ('Numpy 1', 'copia_subjanelas_np1'),
         ('Numpy 2', 'copia_subjanelas_np2'),
         ('Cython', 'copia_subjanelas_cy'),
         ('Numba', 'copia_subjanelas_puro_nb'),
         ('Numba + Numpy', 'copia_subjanelas_nb_np')]

print('''
Função        Min    Max
------------- ------ ------''')
for n, f in funcs:
    values = timeit.repeat('%s(img, msk, win, X)'%f,
globals=globals(), number=nexec, repeat=repeat)
    mn = min(values)
    mx = max(values)
    print('%s %0.2f %0.2f'%(n.ljust(13), mn, mx))

print('')
Número de subjanelas copiadas 225600

Função        Min    Max
------------- ------ ------
Python        85.30 91.36
Numpy 1       16.72 18.63
Numpy 2       10.86 13.77
Cython        0.34 0.39
Numba         0.32 0.58
Numba + Numpy 0.26 0.59

Novamente a versão em Python é mais lenta por várias ordens de magnitude. A utilização de np.nonzeros realmente deixa o processamento mais rápido, mas a principal diferença entre as versões usando Numpy a o Python puro é a cópia da subimagem. Somente esta modificação diminui o tempo em mais de 4 vezes. Porém, a versão em Numpy ainda está muito longe das otimizadas usando Cython e Numba. Ambos aceleradores tiveram desempenho mínimo parecido, apesar de Numba possuir um tempo de execução lento da primeira vez.

Conclusão

Apesar dos testes serem simples, podemos ver uma certa ordem de prioridade ao otimizar código Python. Primeiramente, nunca escreva código em Python puro para processamento de dados. Utilizar ao menos Numpy é obrigatório e já tornará o código muito mais rápido. Toda computação que puder ser feita usando expressões em Numpy deve ser feita desta maneira. Segundo, a velocidade obtida por código Numba e Cython é muito próxima. Cython é muito mais flexível e permite decidir a compatibilidade de uma expressão em vez de uma função inteira. Por outro lado, Numba pode ser tão rápido quanto Cython para um pequeno subconjunto de programas em Python, mas não é flexível.

Pessoalmente, tenho tendido cada vez mais a explorar a utilização de Numba ao invés de Cython. O suporte de Numba a Windows é melhor (se usamos Anaconda Python, é claro) e o código costuma ser mais parecido com Python puro.

Nenhum comentário:

Postar um comentário