Google I/O สำเร็จแล้ว! ติดตามเซสชัน TensorFlow ดูเซสชัน

ตุ้มน้ำหนักแบบเบาบางโดยใช้การตัดแต่งกิ่งโครงสร้าง

ดูบน TensorFlow.org ทำงานใน Google Colab ดูแหล่งที่มาบน GitHub ดาวน์โหลดโน๊ตบุ๊ค

ตุ้มน้ำหนักการตัดแต่งโครงสร้างจากแบบจำลองของคุณเพื่อให้กระจัดกระจายในรูปแบบเฉพาะสามารถเร่งเวลาการอนุมานแบบจำลองได้ด้วยการรองรับ HW ที่เหมาะสม

บทช่วยสอนนี้แสดงวิธี:

  • กำหนดและฝึกแบบจำลองบนชุดข้อมูล mnist ด้วยโครงสร้างที่กระจัดกระจาย
  • แปลงโมเดล pruned เป็นรูปแบบ tflite
  • เห็นภาพโครงสร้างของตุ้มน้ำหนักที่ตัดแต่งแล้ว

สำหรับภาพรวมทั่วไปของเทคนิคการตัดแต่งกิ่งเพื่อเพิ่มประสิทธิภาพรูปแบบให้ดูที่ ภาพรวมการตัดแต่งกิ่ง สำหรับการสอนเกี่ยวกับการตัดแต่งกิ่งน้ำหนักทั่วไปดู การตัดแต่งกิ่งใน Keras

การตัดแต่งกิ่งโครงสร้างตุ้มน้ำหนัก

การตัดแต่งกิ่งโครงสร้างจะทำให้น้ำหนักของแบบจำลองเป็นศูนย์ในตอนเริ่มต้นของกระบวนการฝึกอบรม คุณสามารถใช้เทคนิคการตัดแต่งกิ่งนี้บล็อกปกติของน้ำหนักเพื่อเพิ่มความเร็วในการอนุมานในการสนับสนุน HWS ตัวอย่างเช่นการจัดกลุ่มน้ำหนักในแบบจำลองโดยบล็อกสี่และ zeroing ออกมาสองของน้ำหนักผู้ที่อยู่ในแต่ละบล็อกหรือที่เรียกว่าการลดลง 2 4 เทคนิคนี้ใช้กับมิติสุดท้ายของเมตริกซ์น้ำหนักสำหรับรุ่นที่แปลงโดย TensorFlow Lite เท่านั้น ยกตัวอย่างเช่น Conv2D น้ำหนักชั้นใน TensorFlow Lite มีโครงสร้าง [channel_out, height, width, channel_in] และ Dense น้ำหนักชั้นมีโครงสร้าง [channel_out, channel_in] รูปแบบ sparsity ถูกนำไปใช้น้ำหนักในมิติที่ผ่านมา: channel_in

เมื่อเทียบกับการกระจายแบบกระจายแบบสุ่ม โครงสร้างแบบกระจายโดยทั่วไปมีความแม่นยำต่ำกว่าเนื่องจากโครงสร้างที่จำกัด อย่างไรก็ตาม สามารถลดเวลาในการอนุมานได้อย่างมากบนฮาร์ดแวร์ที่รองรับ

การตัดแต่งกิ่งสามารถใช้กับแบบจำลองร่วมกับเทคนิคการบีบอัดแบบจำลองอื่นๆ เพื่อให้อัตราการบีบอัดดีขึ้น ดูควอนและการจัดกลุ่มตัวอย่างใน เทคนิคการเพิ่มประสิทธิภาพการทำงานร่วมกัน สำหรับรายละเอียดเพิ่มเติม

ติดตั้ง

เตรียมสภาพแวดล้อมและข้อมูลการพัฒนาของคุณ

 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

และดาวน์โหลดภาพปกติข้อมูลจาก MNIST ชุด

# 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

กำหนดพารามิเตอร์การตัดแต่งโครงสร้าง

กำหนดพารามิเตอร์สำหรับการตัดแต่งกิ่งและระบุประเภทของการตัดแต่งกิ่งโครงสร้าง ตั้งค่าพารามิเตอร์สำหรับการตัดแต่งกิ่งเพื่อ (2, 4) การตั้งค่าเหล่านี้หมายความว่าในกลุ่มขององค์ประกอบสี่องค์ประกอบ อย่างน้อยสององค์ประกอบที่มีขนาดต่ำสุดจะถูกตั้งค่าเป็นศูนย์

คุณไม่จำเป็นต้องตั้งค่า pruning_schedule พารามิเตอร์ ตามค่าเริ่มต้น รูปแบบการตัดแต่งกิ่งถูกกำหนดไว้ในขั้นตอนแรก และจะไม่มีการปรับปรุงในระหว่างการฝึก

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

กำหนดพารามิเตอร์สำหรับการตัดแต่งกิ่งแบบสุ่มโดยมีเป้าหมายเบาบางเป้าหมาย 50%

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

กำหนดสถาปัตยกรรมแบบจำลองและระบุเลเยอร์ที่จะตัด การตัดแต่งโครงสร้างจะใช้ตามชั้นของแบบจำลองที่คุณเลือก

ในตัวอย่างด้านล่าง เราตัดเฉพาะบางเลเยอร์เท่านั้น เราตัดที่สอง Conv2D ชั้นและเป็นครั้งแรก Dense ชั้น

แจ้งให้ทราบว่าครั้งแรกที่ Conv2D ชั้นไม่สามารถตัดแต่งโครงสร้าง ในการตัดแต่งโครงสร้าง ควรมีมากกว่าหนึ่งช่องสัญญาณเข้า แต่เราตัดครั้งแรก Conv2D ชั้นที่มีการตัดแต่งกิ่งแบบสุ่ม

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
_________________________________________________________________

ฝึกและประเมินแบบจำลอง

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

นำกระดาษห่อหุ้มการตัดแต่งกิ่งออกเพื่อไม่ให้รวมอยู่ในโมเดลเมื่อคุณแปลงเป็นรูปแบบ TensorFlow Lite

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

แปลงโมเดลเป็นรูปแบบ 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

เห็นภาพและตรวจสอบน้ำหนัก

ตอนนี้เห็นภาพโครงสร้างของน้ำหนักใน Dense ชั้น pruned มี 2 4 sparsity แยกน้ำหนักออกจากไฟล์ 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.

ในการตรวจสอบว่าเราเลือกเลเยอร์ที่ถูกต้องซึ่งได้รับการตัดแต่งแล้ว ให้พิมพ์รูปร่างของเทนเซอร์ตุ้มน้ำหนัก

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

ตอนนี้เราเห็นภาพโครงสร้างสำหรับส่วนย่อยของเมตริกซ์น้ำหนัก โครงสร้างของเมตริกซ์น้ำหนักเบาบางในมิติที่ผ่านมาโดยใช้ (2,4) รูปแบบ: สององค์ประกอบในสี่เป็นศูนย์ เพื่อให้การแสดงภาพชัดเจนยิ่งขึ้น เราจึงแทนที่ค่าที่ไม่ใช่ศูนย์ทั้งหมดด้วยค่าต่างๆ

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)

กำหนดฟังก์ชันเสริมเพื่อวาดเส้นแบ่งเพื่อให้เห็นโครงสร้างชัดเจน

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

ทีนี้ลองนึกภาพเซตย่อยของเทนเซอร์น้ำหนัก

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

น้ำหนักเห็นภาพสำหรับ Conv2D ชั้น sparsity โครงสร้างถูกนำไปใช้ในช่องสุดท้ายที่คล้ายกับ Dense ชั้น แค่สอง Conv2D ชั้น pruned โครงสร้างเป็นแหลมออกมาข้างต้น

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

คล้ายกับน้ำหนักของ Dense ชั้นมิติสุดท้ายของเคอร์เนลมี (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

เรามาดูกันว่าตุ้มน้ำหนักที่ตัดแต่งแบบสุ่มเหล่านั้นจะมีลักษณะอย่างไร เราแยกพวกมันและแสดงส่วนย่อยของเมตริกซ์น้ำหนัก

# 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

TensorFlow รุ่นการเพิ่มประสิทธิภาพเครื่องมือรวมถึงหลามสคริปต์ที่สามารถใช้ในการตรวจสอบว่าที่ชั้นในแบบจำลองจากไฟล์ tflite ที่กำหนดมีน้ำหนัก pruned โครงสร้าง: check_sparsity_m_by_n.py คำสั่งต่อไปนี้สาธิตวิธีใช้เครื่องมือนี้เพื่อตรวจสอบความกระจัดกระจายแบบ 2 คูณ 4 ในแบบจำลองเฉพาะ

 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