Google I/O is a wrap! Catch up on TensorFlow sessions View sessions

Pesos esparsos usando poda estrutural

Veja no TensorFlow.org Executar no Google Colab Ver fonte no GitHub Baixar caderno

Os pesos de poda estruturais do seu modelo para torná-lo esparso em um padrão específico podem acelerar o tempo de inferência do modelo com suportes de HW apropriados.

Este tutorial mostra como:

  • Defina e treine um modelo no conjunto de dados mnist com uma dispersão estrutural específica
  • Converta o modelo podado para o formato tflite
  • Visualize a estrutura dos pesos podados

Para uma visão geral da técnica de poda para a otimização do modelo, consulte a visão geral poda . Para tutorial sobre poda de peso geral, consulte Poda na Keras .

Poda estrutural de pesos

A poda estrutural zera sistematicamente os pesos do modelo no início do processo de treinamento. De aplicar esta técnica poda de blocos regulares de pesos para acelerar inferência no apoio HWS, por exemplo: agrupamento pesos no modelo por blocos de quatro e zeragem fora dois desses pesos em cada bloco, conhecido como um 2 4 por redução. Essa técnica se aplica apenas à última dimensão do tensor de peso para o modelo que é convertido pelo TensorFlow Lite. Por exemplo, Conv2D pesos de camada em TensorFlow Lite tem a estrutura [channel_out, height, width, channel_in] e Dense pesos de camada tem a estrutura [channel_out, channel_in] . O padrão de dispersão é aplicado aos pesos na última dimensão: channel_in .

Comparado com a dispersão aleatória, a dispersão estruturada geralmente tem menor precisão devido à estrutura restritiva, no entanto, pode reduzir significativamente o tempo de inferência no hardware suportado.

A poda pode ser aplicada a um modelo junto com outras técnicas de compressão de modelo para uma melhor taxa de compressão. Veja exemplos de quantização e de agrupamento em técnica de otimização colaborativo para mais detalhes.

Configurar

Prepare seu ambiente de desenvolvimento e dados.

 pip install -q tensorflow
 pip install -q tensorflow-model-optimization
 pip install -q matplotlib
import tensorflow as tf
from tensorflow import keras

import tensorflow_model_optimization as tfmot
prune_low_magnitude = tfmot.sparsity.keras.prune_low_magnitude

Download e imagem normalizar dados do MNIST dataset

# Load MNIST dataset.
mnist = keras.datasets.mnist
(train_images, train_labels), (test_images, test_labels) = mnist.load_data()

# Normalize the input image so that each pixel value is between 0 and 1.
train_images = train_images / 255.0
test_images = test_images / 255.0
Downloading data from https://storage.googleapis.com/tensorflow/tf-keras-datasets/mnist.npz
11493376/11490434 [==============================] - 0s 0us/step
11501568/11490434 [==============================] - 0s 0us/step

Definir parâmetros de poda estrutural

Defina parâmetros para poda e especifique o tipo de poda estrutural. Definir os parâmetros para a poda de (2, 4) . Essas configurações significam que em um bloco de quatro elementos, pelo menos dois com a menor magnitude são zerados.

Você não tem que definir o pruning_schedule parâmetro. Por padrão, a máscara de poda é definida na primeira etapa e não é atualizada durante o treinamento.

pruning_params_2_by_4 = {
    'sparsity_m_by_n': (2, 4),
}

Defina parâmetros para poda aleatória com a esparsidade alvo de 50%.

pruning_params_sparsity_0_5 = {
    'pruning_schedule': tfmot.sparsity.keras.ConstantSparsity(target_sparsity=0.5,
                                                              begin_step=0,
                                                              frequency=100)
}

Defina a arquitetura do modelo e especifique quais camadas devem ser podadas. A poda estrutural é aplicada com base nas camadas do modelo selecionado.

No exemplo abaixo, podamos apenas algumas das camadas. Nós podar segundo Conv2D camada e o primeiro Dense camada.

Note-se que o primeiro Conv2D camada não pode ser podada estruturalmente. Para ser podado estruturalmente, deve ter mais de um canal de entrada. Em vez disso, podar o primeiro Conv2D camada com poda aleatória.

model = keras.Sequential([
    prune_low_magnitude(
        keras.layers.Conv2D(
            32, 5, padding='same', activation='relu',
            input_shape=(28, 28, 1),
            name="pruning_sparsity_0_5"),
        **pruning_params_sparsity_0_5),
    keras.layers.MaxPooling2D((2, 2), (2, 2), padding='same'),
    prune_low_magnitude(
        keras.layers.Conv2D(
            64, 5, padding='same',
            name="structural_pruning"),
        **pruning_params_2_by_4),
    keras.layers.BatchNormalization(),
    keras.layers.ReLU(),
    keras.layers.MaxPooling2D((2, 2), (2, 2), padding='same'),
    keras.layers.Flatten(),
    prune_low_magnitude(
        keras.layers.Dense(
            1024, activation='relu',
            name="structural_pruning_dense"),
        **pruning_params_2_by_4),
    keras.layers.Dropout(0.4),
    keras.layers.Dense(10)
])

model.compile(optimizer='adam',
              loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
              metrics=['accuracy'])

model.summary()
2022-01-18 12:20:52.816023: E tensorflow/stream_executor/cuda/cuda_driver.cc:271] failed call to cuInit: CUDA_ERROR_NO_DEVICE: no CUDA-capable device is detected
/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow_model_optimization/python/core/sparsity/keras/pruning_wrapper.py:218: UserWarning: `layer.add_variable` is deprecated and will be removed in a future version. Please use `layer.add_weight` method instead.
  aggregation=tf.VariableAggregation.MEAN)
/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow_model_optimization/python/core/sparsity/keras/pruning_wrapper.py:225: UserWarning: `layer.add_variable` is deprecated and will be removed in a future version. Please use `layer.add_weight` method instead.
  aggregation=tf.VariableAggregation.MEAN)
/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow_model_optimization/python/core/sparsity/keras/pruning_wrapper.py:238: UserWarning: `layer.add_variable` is deprecated and will be removed in a future version. Please use `layer.add_weight` method instead.
  trainable=False)
Model: "sequential"
_________________________________________________________________
 Layer (type)                Output Shape              Param #   
=================================================================
 prune_low_magnitude_pruning  (None, 28, 28, 32)       1634      
 _sparsity_0_5 (PruneLowMagn                                     
 itude)                                                          
                                                                 
 max_pooling2d (MaxPooling2D  (None, 14, 14, 32)       0         
 )                                                               
                                                                 
 prune_low_magnitude_structu  (None, 14, 14, 64)       102466    
 ral_pruning (PruneLowMagnit                                     
 ude)                                                            
                                                                 
 batch_normalization (BatchN  (None, 14, 14, 64)       256       
 ormalization)                                                   
                                                                 
 re_lu (ReLU)                (None, 14, 14, 64)        0         
                                                                 
 max_pooling2d_1 (MaxPooling  (None, 7, 7, 64)         0         
 2D)                                                             
                                                                 
 flatten (Flatten)           (None, 3136)              0         
                                                                 
 prune_low_magnitude_structu  (None, 1024)             6423554   
 ral_pruning_dense (PruneLow                                     
 Magnitude)                                                      
                                                                 
 dropout (Dropout)           (None, 1024)              0         
                                                                 
 dense (Dense)               (None, 10)                10250     
                                                                 
=================================================================
Total params: 6,538,160
Trainable params: 3,274,762
Non-trainable params: 3,263,398
_________________________________________________________________

Treinar e avaliar o modelo.

batch_size = 128
epochs = 2

model.fit(
    train_images,
    train_labels,
    batch_size=batch_size,
    epochs=epochs,
    verbose=0,
    callbacks=tfmot.sparsity.keras.UpdatePruningStep(),
    validation_split=0.1)

_, pruned_model_accuracy = model.evaluate(test_images, test_labels, verbose=0)
print('Pruned test accuracy:', pruned_model_accuracy)
Pruned test accuracy: 0.9861000180244446

Remova o wrapper de poda para que ele não seja incluído no modelo ao convertê-lo para o formato TensorFlow Lite.

model = tfmot.sparsity.keras.strip_pruning(model)

Converter modelo para formato tflite

import tempfile

converter = tf.lite.TFLiteConverter.from_keras_model(model)
tflite_model = converter.convert()

_, tflite_file = tempfile.mkstemp('.tflite')
print('Saved converted pruned model to:', tflite_file)
with open(tflite_file, 'wb') as f:
  f.write(tflite_model)
2022-01-18 12:21:27.701505: W tensorflow/python/util/util.cc:368] Sets are not currently considered sequences, but this may change in the future, so consider avoiding using them.
INFO:tensorflow:Assets written to: /tmp/tmpcg_uiteu/assets
2022-01-18 12:21:28.552924: W tensorflow/compiler/mlir/lite/python/tf_tfl_flatbuffer_helpers.cc:357] Ignored output_format.
2022-01-18 12:21:28.552974: W tensorflow/compiler/mlir/lite/python/tf_tfl_flatbuffer_helpers.cc:360] Ignored drop_control_dependency.
WARNING:absl:Buffer deduplication procedure will be skipped when flatbuffer library is not properly loaded
Saved converted pruned model to: /tmp/tmp0w1w7jvg.tflite

Visualize e verifique os pesos

Agora visualizar a estrutura de pesos na Dense camada podadas com 2 por 4 esparsidade. Extraia os pesos do arquivo tflite.

# Load tflite file with the created pruned model
interpreter = tf.lite.Interpreter(model_path=tflite_file)
interpreter.allocate_tensors()

details = interpreter.get_tensor_details()

# Weights of the dense layer that has been pruned.
tensor_name = 'structural_pruning_dense/MatMul'
detail = [x for x in details if tensor_name in x["name"]]

# We need the first layer.
tensor_data = interpreter.tensor(detail[0]["index"])()
INFO: Created TensorFlow Lite XNNPACK delegate for CPU.

Para verificar se selecionamos a camada correta que foi podada, imprima a forma do tensor de peso.

print(f"Shape of Dense layer is {tensor_data.shape}")
Shape of Dense layer is (1024, 3136)

Agora visualizamos a estrutura para um pequeno subconjunto do tensor de peso. A estrutura do tensor de peso é escasso na última dimensão, utilizando o (2,4) padrão: dois elementos de cada quatro são zeros. Para tornar a visualização mais clara, substituímos todos os valores diferentes de zero por uns.

import matplotlib.pyplot as plt
import numpy as np

# The value 24 is chosen for convenience.
width = height = 24

subset_values_to_display = tensor_data[0:height, 0:width]

val_ones = np.ones([height, width])
val_zeros = np.zeros([height, width])
subset_values_to_display = np.where(abs(subset_values_to_display) > 0, val_ones, val_zeros)

Defina a função auxiliar para desenhar linhas de separação para ver a estrutura claramente.

def plot_separation_lines(height, width):

    block_size = [1, 4]

    # Add separation lines to the figure.
    num_hlines = int((height - 1) / block_size[0])
    num_vlines = int((width - 1) / block_size[1])
    line_y_pos = [y * block_size[0] for y in range(1, num_hlines + 1)]
    line_x_pos = [x * block_size[1] for x in range(1, num_vlines + 1)]

    for y_pos in line_y_pos:
        plt.plot([-0.5, width], [y_pos - 0.5 , y_pos - 0.5], color='w')

    for x_pos in line_x_pos:
        plt.plot([x_pos - 0.5, x_pos - 0.5], [-0.5, height], color='w')

Agora visualize o subconjunto do tensor de peso.

plot_separation_lines(height, width)

plt.axis('off')
plt.imshow(subset_values_to_display)
plt.colorbar()
plt.title("Structural pruning for Dense layer")
plt.show()

png

Visualizar pesos para o Conv2D camada. A dispersão estrutural é aplicada no último canal, semelhante à Dense camada. Apenas o segundo Conv2D camada é estruturalmente podadas como fora pontiagudo acima.

# Get weights of the convolutional layer that has been pruned with 2 by 4 sparsity.
tensor_name = 'structural_pruning/Conv2D'
detail = [x for x in details if tensor_name in x["name"]]
tensor_data = interpreter.tensor(detail[1]["index"])()
print(f"Shape of the weight tensor is {tensor_data.shape}")
Shape of the weight tensor is (64, 5, 5, 32)

Semelhante aos pesos dos Dense camada, a última dimensão do núcleo tem uma estrutura (2, 4).

weights_to_display = tf.reshape(tensor_data, [tf.reduce_prod(tensor_data.shape[:-1]), -1])
weights_to_display = weights_to_display[0:width, 0:height]

val_ones = np.ones([height, width])
val_zeros = np.zeros([height, width])
subset_values_to_display = np.where(abs(weights_to_display) > 1e-9, val_ones, val_zeros)

plot_separation_lines(height, width)

plt.axis('off')
plt.imshow(subset_values_to_display)
plt.colorbar()
plt.title("Structurally pruned weights for Conv2D layer")
plt.show()

png

Vamos ver como ficam esses pesos podados aleatoriamente. Nós os extraímos e exibimos um subconjunto do tensor de peso.

# Get weights of the convolutional layer that has been pruned with random pruning.
tensor_name = 'pruning_sparsity_0_5/Conv2D'
detail = [x for x in details if tensor_name in x["name"]]
tensor_data = interpreter.tensor(detail[0]["index"])()
print(f"Shape of the weight tensor is {tensor_data.shape}")
Shape of the weight tensor is (32, 5, 5, 1)
weights_to_display = tf.reshape(tensor_data, [tensor_data.shape[0],tf.reduce_prod(tensor_data.shape[1:])])
weights_to_display = weights_to_display[0:width, 0:height]

val_ones = np.ones([height, width])
val_zeros = np.zeros([height, width])
subset_values_to_display = np.where(abs(weights_to_display) > 0, val_ones, val_zeros)

plot_separation_lines(height, width)

plt.axis('off')
plt.imshow(subset_values_to_display)
plt.colorbar()
plt.title("Unstructed pruned weights for Conv2D layer")
plt.show()

png

O TensorFlow modelo de otimização Toolkit inclui um script python que pode ser usado para verificar se que camadas do modelo a partir do determinado arquivo tflite têm os pesos estruturalmente podadas: check_sparsity_m_by_n.py . O comando a seguir demonstra como usar essa ferramenta para verificar a esparsidade de 2 por 4 em um modelo específico.

 python3 ./tensorflow_model_optimization/python/core/sparsity/keras/tools/check_sparsity_m_by_n.py --model_tflite=pruned_model.tflite --m_by_n=2,4
python3: can't open file './tensorflow_model_optimization/python/core/sparsity/keras/tools/check_sparsity_m_by_n.py': [Errno 2] No such file or directory