Google I / O là một kết quả hoàn hảo! Cập nhật các phiên TensorFlow Xem phiên

Trọng lượng thưa sử dụng cắt tỉa cấu trúc, Trọng lượng thưa sử dụng cắt tỉa cấu trúc

Xem trên TensorFlow.org Chạy trong Google Colab Xem nguồn trên GitHub Tải xuống sổ ghi chép

Việc cắt tỉa cấu trúc trọng số từ mô hình của bạn để làm cho nó thưa thớt trong mô hình cụ thể có thể đẩy nhanh thời gian suy luận mô hình với các hỗ trợ HW thích hợp.

Hướng dẫn này chỉ cho bạn cách:

  • Xác định và đào tạo một mô hình trên tập dữ liệu mnist với một cấu trúc thưa thớt cụ thể
  • Chuyển đổi mô hình đã được lược bớt sang định dạng tflite
  • Hình dung cấu trúc của các quả cân được cắt tỉa

Để có cái nhìn tổng quát về kỹ thuật tỉa cho việc tối ưu hóa mô hình, xem tổng quan tỉa . Đối với hướng dẫn trên tỉa trọng lượng nói chung, xem Tỉa trong Keras .

Kết cấu của tạ

Việc cắt tỉa cấu trúc một cách có hệ thống giúp loại bỏ các trọng số của mô hình khi bắt đầu quá trình đào tạo. Bạn áp dụng điều này các kỹ thuật cắt tỉa để khối thường xuyên của trọng lượng để tăng tốc độ suy luận vào việc hỗ trợ CTNH, ví dụ: nhóm trọng trong mô hình của khối bốn và zeroing ra hai trong số những trọng trong mỗi khối, được gọi là một 2 của 4 giảm. Kỹ thuật này chỉ áp dụng cho kích thước cuối cùng của bộ căng cân đối với mô hình được chuyển đổi bởi TensorFlow Lite. Ví dụ, Conv2D trọng lượng lớp trong TensorFlow Lite có cấu trúc [channel_out, height, width, channel_in]Dense trọng lượng lớp có cấu trúc [channel_out, channel_in] . Các mô hình thưa thớt được áp dụng cho các trọng trong chiều cuối: channel_in .

So với độ thưa thớt ngẫu nhiên, độ thưa thớt có cấu trúc thường có độ chính xác thấp hơn do cấu trúc hạn chế, tuy nhiên, nó có thể giảm đáng kể thời gian suy luận trên phần cứng được hỗ trợ.

Việc cắt tỉa có thể được áp dụng cho một mô hình cùng với các kỹ thuật nén mô hình khác để có tốc độ nén tốt hơn. Xem ví dụ về lượng tử và clustering trong kỹ thuật tối ưu hóa hợp tác để biết thêm chi tiết.

Thành lập

Chuẩn bị môi trường phát triển và dữ liệu của bạn.

 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

Tải về và hình ảnh bình thường hóa dữ liệu từ MNIST bộ dữ liệu

# 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

Xác định các thông số cắt tỉa cấu trúc

Xác định các thông số để cắt tỉa và chỉ định kiểu cắt tỉa cấu trúc. Thiết lập thông số cho cắt tỉa để (2, 4) . Các cài đặt này có nghĩa là trong một khối gồm bốn phần tử, ít nhất hai phần tử có độ lớn thấp nhất được đặt thành không.

Bạn không cần phải thiết lập các pruning_schedule tham số. Theo mặc định, mặt nạ cắt tỉa được xác định ở bước đầu tiên và nó không được cập nhật trong quá trình đào tạo.

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

Xác định các thông số để cắt tỉa ngẫu nhiên với độ thưa thớt mục tiêu là 50%.

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

Xác định kiến ​​trúc mô hình và chỉ định các lớp cần cắt bớt. Việc cắt tỉa cấu trúc được áp dụng dựa trên các lớp của mô hình bạn chọn.

Trong ví dụ dưới đây, chúng tôi chỉ cắt bớt một số lớp. Chúng tôi prune thứ hai Conv2D lớp và là người đầu tiên Dense lớp.

Chú ý rằng người đầu tiên Conv2D lớp không thể được bỏ bớt cấu trúc. Để được lược bớt về mặt cấu trúc, nó nên có nhiều hơn một kênh đầu vào. Thay vào đó, chúng tôi prune là người đầu tiên Conv2D lớp với tỉa ngẫu nhiên.

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
_________________________________________________________________

Huấn luyện và đánh giá mô hình.

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

Loại bỏ trình bao bọc cắt tỉa để nó không có trong mô hình khi bạn chuyển đổi nó sang định dạng TensorFlow Lite.

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

Chuyển đổi mô hình sang định dạng 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

Hình dung và kiểm tra trọng lượng

Bây giờ hình dung cấu trúc của trọng trong Dense lớp cắt tỉa với 2 bằng 4 thưa thớt. Trích xuất các trọng số từ tệp 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.

Để xác minh rằng chúng tôi đã chọn đúng lớp đã được cắt bớt, hãy in hình dạng của bộ căng khối lượng.

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

Bây giờ chúng ta hình dung cấu trúc của một tập con nhỏ của tensor trọng lượng. Cấu trúc của tensor trọng lượng là thưa thớt trong chiều qua, bằng cách sử dụng (2,4) mẫu: hai yếu tố trong số bốn là số không. Để hiển thị rõ ràng hơn, chúng tôi thay thế tất cả các giá trị khác 0 bằng các giá trị khác.

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)

Xác định chức năng phụ vẽ đường phân cách để nhìn rõ kết cấu.

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')

Bây giờ hình dung tập hợp con của tensor trọng lượng.

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

Hình dung trọng cho Conv2D lớp. Các thưa thớt cấu trúc được áp dụng trong các kênh cuối cùng, tương tự như Dense lớp. Chỉ có thứ hai Conv2D lớp được cấu trúc tỉa như ra nhọn ở trên.

# 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)

Tương tự như trọng lượng của Dense lớp, kích thước cuối cùng của hạt nhân có một (2, 4) cấu trúc.

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

Hãy xem những quả cân được cắt tỉa ngẫu nhiên đó trông như thế nào. Chúng tôi giải nén chúng và hiển thị một tập hợp con của tensor trọng lượng.

# 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

Các TensorFlow Mẫu Tối ưu hóa Toolkit bao gồm một kịch bản python mà có thể được sử dụng để kiểm tra xem có các lớp trong mô hình từ tập tin tflite cho có trọng lượng cấu trúc tỉa: check_sparsity_m_by_n.py . Lệnh sau đây trình bày cách sử dụng công cụ này để kiểm tra độ thưa thớt 2 x 4 trong một mô hình cụ thể.

 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