Medindo consumo de memória em Python

O gerenciamento de memória de programas em Python é feito inteiramente pela linguagem. Isto significa que, por vezes, pode ser difícil controlar o uso de memória e não costuma ser claro o quanto cada objeto ocupa na memória nem quando seu espaço é liberado. Descreverei neste post o uso do módulo memory_profiler, que contém ferramentas para medir o consumo de memória de scripts Python.

Antes de tudo, o memory_profiler pode ser instalado via pip usando o seguinte comando. A opção -U pode ser usada para instalar o módulo na área do usuário em sistemas Unix-like.

$ pip install memory_profiler

Para usar este módulo basta decorar a função que desejamos medir o consumo de memória com o decorador @profile e executar o programa usando

$ python -m memory_profiler arquivo.py

Este comando imprime um relatório linha a linha mostrando o consumo de memória. Quando objetos são criados o consumo de memória aumenta, enquanto quando eles saem de contexto o consumo pode diminuir. O ponto fraco deste modo de execução é que ele só mostra o consumo medido "entre" as linhas do programa. Ou seja, se uma função for chamada e, durante sua execução alocar e liberar uma grande quantidade de memória estes valores não são contabilizados.

Também é disponibilizado o comando mprof, que executa um script e mede o consumo total de memória em pequenos intervalos de tempo. mprof, porém, não discrima onde a memória está sendo usada.

---

Irei exemplificar o uso do mprof com dois exemplos tirados da TRIOSlib, a biblioteca que desenvolvi durante meu doutorado. O objetivo da biblioteca é usar aprendizado de máquina para tratar problemas de processamento de imagens (mais detalhes aqui).

Uma das tarefas que mais consome memória na TRIOSlib é a extração de características. Um padrão é extraído de cada ponto das imagens de treinamento. Conjuntos de treinamento contendo várias centenas de milhares de padrões são comuns, mesmo usando um pequeno número de imagens de treinamento.

A TRIOSlib implementa dois modos para extração de características. No primeiro uma matriz é alocada com uma linha por pixel das imagens de treinamento. Cada padrão observado é armazenado em uma linha e, se houverem repetições, o mesmo padrão é armazenado diversas vezes. Neste caso o consumo depende do número de pixels das imagens de treinamento. Nos referimos a este modo como modo Array.

No segundo modo a matriz é trocada por um dicionário cujas chaves são padrões observados nas imagens. Se houverem repetições não há consumo extra de memória. Neste caso o consumo de memória depende do número de padrões únicos nas imagens. Quanto maior a repetição de padrões menor será o consumo relativo ao primeiro modo. Nos referimos a este tipo de execução como modo Dicionário ou Dict.

Apesar do memory_profiler possuir um modo que mostra o consumo de memória linha a linha, para operações longas como o treinamento de operadores pode ser mais interessante usar o comando mprof. Apesar de não ser possível identificar qual estrutura está consumindo memória, podemos ver o quanto a memória cresce e diminui conforme o programa executa. Também conseguimos, por meio do decorador @profile, marcar o início e o fim da execução de funções no código.

Para executar um script usamos o comando mprof run arquivo.py. Podemos plotar um gráfico com a última execução usando mprof plot. Usaremos o código abaixo como exemplo.

import trios
import trios.feature_extractors
import trios.classifiers

import numpy as np
from sklearn.tree import DecisionTreeClassifier

import sys


@profile
def main(ordered=True):
    imgset =
trios.Imageset.read('/media/igor/Data1/datasets/staffs/level1.set')
    testset =
trios.Imageset.read('/media/igor/Data1/datasets/staffs/test-
small.set')
    win = np.ones((7, 7), np.uint8)

    wop = trios.WOperator(win,
trios.classifiers.SKClassifier(DecisionTreeClassifier(),
ordered=ordered), trios.feature_extractors.RAWFeatureExtractor(win))

    dataset = wop.extractor.extract_dataset(imgset, ordered=ordered)
    if ordered:
        print(dataset[1].shape[0])
    else:
        print(len(dataset))

    tr2 = profile(wop.classifier.train)
    tr2(dataset, {})
    wop.trained = True

    print(wop.eval(testset[:1]))

if __name__ == '__main__':
    main(sys.argv[1] == 'ordered')

O decorator @profile irá marcar no gráfico o início e fim das funções main e WOperator.Classifier.train. (Lembre-se que um decorador é simplesmente uma função que retorna outra função. Fazemos isto na linha 21 para obter uma função de treinamento com @profile.) A escolha do modo de extração de características é feita pelo parâmetro ordered, passado pela linha de comando. ordered=True implica no uso de uma matriz para os padrões. ordered=False usa um dicionário. Veja abaixo o consumo de memória medidos em uma execução usando os dois modos.

Fig 1: Modo Array

Fig 2: Modo Dict

O consumo máximo de memória do modo Array é cerca de 4 vezes maior do que o modo Dict! A principal diferença entre os dois modos de execução é a quantidade de memória usada pelo método train. Enquanto o modo Array executa a extração de características muito mais rápido que o modo Dict, o treinamento é muito mais custoso tanto em termos de tempo quanto em memória. Isto pode ser explicado pelo fato de que o modo Dict condensa todas as observações repetidas em uma só linha ao treinar. Isto resulta em uma matriz de treinamento com 137.453 linhas, contra 5.595.585 linhas quando usamos o modo Array. A redução no número de linhas é possível atribuindo ao peso de cada exemplo de treinamento único o número de repetições observadas nas imagens. Infelizmente a maioria dos classificadores do scikit-learn não oferece suporte a definir pesos para cada instância de treinamento. Classificadores baseados em árvores de decisão (encontrados no pacote sklearn.tree) possuem este suporte e, como podemos ver isto faz uma grande diferença ao trabalhar com dados em que há repetição nos padrões de entrada observados.

---

Neste texto exemplifiquei a utilização do memory_profiler para medição de consumo de memória em Python e do decorador @profile para marcar o início e o fim da execução de diversas funções. Como pudemos ver, utilizar estruturas adequadas pode diminuir significativamente o consumo de memória e permitir processar uma maior quantidade de dados em memória.