млн операций в секунду
Оптимизация модели с помощью TensorFlow
Сократите задержку, хранение и затраты на вывод моделей с помощью квантования и обрезки.
За последние несколько лет в моделях машинного обучения наблюдались две, казалось бы, противоположные тенденции. С одной стороны, модели имеют тенденцию становиться все больше и больше, достигая кульминации в том, что сейчас в моде: большие языковые модели. Модель Nvidia Megatron-Turing Natural Language Generation имеет 530 миллиардов параметров! С другой стороны, эти модели развертываются на все более мелких устройствах, таких как смарт-часы или дроны, память и вычислительная мощность которых, естественно, ограничены их размером.
Как втиснуть все более крупные модели во все более компактные устройства? Ответ — оптимизация модели: процесс сжатия модели по размеру и уменьшения ее задержки. В этой статье мы увидим, как это работает и как реализовать два популярных метода оптимизации модели — квантование и обрезку — в TensorFlow.
Базовая модель
Прежде чем мы перейдем к методам оптимизации модели, нам нужно оптимизировать игрушечную модель. Давайте обучим простой бинарный классификатор различать две знаменитые достопримечательности Парижа: Эйфелеву башню и Мону Лизу, нарисованные игроками в игру Google под названием «Быстро рисуй!». Набор данных QuickDraw состоит из изображений в градациях серого размером 28x28.
Давайте обучим простую сверточную сеть для классификации двух ориентиров.
def get_model():
return tf.keras.Sequential([
tf.keras.layers.InputLayer(input_shape=(28, 28)),
tf.keras.layers.Reshape(target_shape=(28, 28, 1)),
tf.keras.layers.Conv2D(
filters=12, kernel_size=(3, 3), activation="relu"
),
tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),
tf.keras.layers.Conv2D(
filters=24, kernel_size=(3, 3), activation="relu"
),
tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),
tf.keras.layers.Flatten(),
tf.keras.layers.Dense(1, activation="sigmoid")
])
model_baseline = get_model()
model_baseline.compile(
optimizer="adam",
loss="binary_crossentropy",
metrics=["accuracy"]
)
model_baseline.fit(
x_train,
y_train,
epochs=1,
)
Теперь мы сохраним модель в формате TensorFlow Lite. Это меньший и более эффективный формат файла по сравнению с традиционным файлом .h5, разработанный специально для мобильных и периферийных развертываний.
converter = tf.lite.TFLiteConverter.from_keras_model(model_baseline)
model_baseline_tflite = converter.convert()
with open("model_baseline.tflite", "wb") as f:
f.write(model_baseline_tflite)
Чтобы оценить размер и точность нашей модели, нам понадобятся две простые функции полезности. Во-первых, evaluate_tflite_model() настраивает интерпретатор TF Lite для передачи тестовых примеров в сохраненную модель TF Lite и вычисляет точность своих прогнозов. Во-вторых, get_gzipped_model_size() создает временную заархивированную версию файла модели .tflite для дальнейшего сжатия и возвращает ее размер на диске в байтах. Реализации обеих функций, которые мы здесь используем, основаны на аналогичных утилитах из Coursera’a Машинное обучение, моделирование конвейеров в производстве.
def evaluate_tflite_model(filename, x_test, y_test):
interpreter = tf.lite.Interpreter(model_path=filename)
interpreter.allocate_tensors()
input_index = interpreter.get_input_details()[0]["index"]
output_index = interpreter.get_output_details()[0]["index"]
y_pred = []
for test_image in x_test:
test_image = np.expand_dims(test_image, axis=0).astype(np.float32)
interpreter.set_tensor(input_index, test_image)
interpreter.invoke()
output = interpreter.tensor(output_index)
y_pred.append(output()[0][0] >= 0.5)
return (y_pred == np.array(y_test)).mean()
def get_gzipped_model_size(file):
_, zipped_file = tempfile.mkstemp(".zip")
with zipfile.ZipFile(
zipped_file, "w", compression=zipfile.ZIP_DEFLATED
) as f:
f.write(file)
return os.path.getsize(zipped_file)
Давайте получим некоторые базовые показатели из нашей неоптимизированной модели.
model_baseline_acc = evaluate_tflite_model(
"model_baseline.tflite", x_test, y_test
)
model_baseline_size = get_gzipped_model_size("model_baseline.tflite")
print(f"Baseline accuracy: {model_baseline_acc}")
print(f"Baseline size: {model_baseline_size}")
Baseline accuracy: 0.9852449095827994
Baseline size: 14303
Модель имеет точность 98,524%, занимая более 14 КБ пространства. Давайте посмотрим, сможем ли мы уменьшить его размер, не жертвуя большей частью точности.
Квантование
Первый метод оптимизации, который мы рассмотрим, — это квантование модели. Его цель — снизить точность объектов, в которых хранятся параметры модели.
Квантование снижает точность объектов, в которых хранятся параметры модели.
По умолчанию TensorFlow сохраняет смещения, веса и активации модели в виде 32-битных чисел с плавающей запятой. Быстрая проверка с помощью np.finfo(np.float32).max покажет вам, что этот тип данных позволяет хранить значения размером до 3³⁸. Но нужно ли это нам?
Как правило, вес и активация большинства моделей не так уж далеки от нуля; в противном случае градиенты взорвутся, не давая нам в первую очередь обучить модель. Преобразование этих float32 в более легкую структуру данных, такую как int8, может иметь большое значение для уменьшения размера модели, не обязательно влияя на ее точность.
Квантование очень удобно использовать, так как оно оперирует уже обученной моделью и преобразует только ее внутренние структуры данных. В TensorFlow это можно сделать при конвертации модели в формат TF Lite, установив атрибут optimizations в конвертере.
converter = tf.lite.TFLiteConverter.from_keras_model(model_baseline)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
model_quantized_tflite = converter.convert()
with open("model_quantized.tflite", "wb") as f:
f.write(model_quantized_tflite)
Это реализует вариант квантования по умолчанию, называемый квантованием динамического диапазона. В большинстве случаев это хорошая отправная точка. Доступны и другие, более продвинутые стратегии, такие как полное целочисленное квантование, которое пытается оценить диапазон значений в тензоре каждой модели, но для этого требуется репрезентативный набор данных для калибровки его оценок.
Давайте остановимся на варианте по умолчанию и проверим размер и точность квантованной модели.
model_quantized_acc = evaluate_tflite_model(
"model_quantized.tflite", x_test, y_test
)
model_quantized_size = get_gzipped_model_size("model_quantized.tflite")
print(f"Quantized accuracy: {model_quantized_acc}")
print(f"Quantized size: {model_quantized_size}")
Quantized accuracy: 0.98526270824434
Quantized size: 7483
Размер модели был уменьшен почти вдвое, с более чем 14 КБ до более 7 КБ! При этом точность не ухудшилась. На самом деле, он даже немного вырос с 98,524% до 98,526%. Это редкий случай, поскольку обычно наблюдается небольшое снижение точности в результате квантования.
Обрезка
Другой популярный метод оптимизации модели — сокращение веса. В больших нейронных сетях очень маловероятно, что каждый вес играет решающую роль в производительности модели. Выявление не столь важных весов и удаление их из сети позволяет нам уменьшить пространство, необходимое для хранения модели, и сэкономить на умножениях с плавающей запятой во время вывода, уменьшая его задержку.
Обрезка определяет и удаляет из модели неважные веса, уменьшая ее размер и время вывода.
Самый простой способ определить неважные веса — по их величине. Веса со значениями, близкими к нулю, с наименьшей вероятностью внесут большой вклад в выход сети и являются лучшими кандидатами на отсечение.
К сожалению, установка некоторых весов на ноль в обученной сети нарушает поток информации внутри нее, что приводит к резкому падению точности. Очевидным решением было бы повторное обучение модели после обрезки, но это оказалось сложной задачей.
Вместо этого обрезку и переобучение можно выполнять одновременно. Это позволяет нам явно научить модель не использовать некоторые веса сети.
Чтобы реализовать сокращение при переобучении в TensorFlow, нам нужно определить расписание сокращения. Расписание регулирует два процесса при переподготовке:
- Какую разреженность применять на протяжении всего обучения;
- Когда (на каких этапах обучения) применять его.
Популярным выбором скорости увеличения разреженности, который доказал свою эффективность на практике, является полиномиальная функция затухания. Это резко увеличит разреженность на первых этапах обучения, чтобы замедлить увеличение на более поздних этапах.
Здесь мы начнем с требования, чтобы 50% весов были равны нулю в начале тренировки, и этот процент увеличится до 80% в конце. Мы также начнем вводить разреженность сразу же, как только начнется обучение, и процесс будет длиться на всех этапах обучения.
batch_size = 128
epochs = 2
validation_split = 0.1
x_train_size = x_train.shape[0] * (1 - validation_split)
end_step = np.ceil(x_train_size / batch_size).astype(np.int32) * epochs
pruning_params = {
"pruning_schedule": tfmot.sparsity.keras.PolynomialDecay(
initial_sparsity=0.50,
final_sparsity=0.80,
begin_step=0,
end_step=end_step,
)
}
Настроив расписание, мы передаем его вместе с обученной базовой моделью в метод prune_low_magnitude. Затем мы перекомпилируем и переобучим модель, передав ей обратный вызов UpdatePruningStep. Этот обратный вызов вызывается после обработки каждого пакета обучающих данных и обновляет шаг сокращения, чтобы он соответствовал заранее определенному расписанию. Он также добавляет обертки вокруг слоев, которые нужно обрезать, поэтому, чтобы получить исходную архитектуру модели, нам нужно удалить эти обертки после повторного обучения с помощью метода strip_pruning.
model_pruned = tfmot.sparsity.keras.prune_low_magnitude(
model_baseline,
**pruning_params,
)
model_pruned.compile(
optimizer="adam",
loss="binary_crossentropy",
metrics=["accuracy"]
)
model_pruned.fit(
x_train,
y_train,
epochs=epochs,
validation_split=validation_split,
callbacks=[tfmot.sparsity.keras.UpdatePruningStep()],
)
model_pruned = tfmot.sparsity.keras.strip_pruning(model_pruned)
Теперь мы можем преобразовать модель в формат TFLite и проверить ее размер и точность, как мы делали это раньше.
converter = tf.lite.TFLiteConverter.from_keras_model(model_pruned)
model_pruned_tflite = converter.convert()
with open("model_pruned.tflite", "wb") as f:
f.write(model_pruned_tflite)
model_pruned_acc = evaluate_tflite_model(
"model_pruned.tflite", x_test, y_test
)
model_pruned_size = get_gzipped_model_size("model_pruned.tflite")
print(f"Pruned accuracy: {model_pruned_acc}")
print(f"Pruned size: {model_pruned_size}")
Pruned accuracy: 0.9840168019364943
Pruned size: 5862
Обрезка сжала нашу модель даже больше, чем квантование: с 14 КБ до менее чем 6 КБ. Однако точность немного пострадала; он упал с 98,524% в базовой модели до 98,402% в сокращенной версии.
Квантование и обрезка объединяются
Поскольку квантование и обрезка работают независимо и по-разному, ничто не запрещает нам использовать оба этих метода одновременно. Мы можем обрезать некоторые веса модели, а затем квантовать оставшиеся для достижения еще более сильного сжатия.
У нас уже есть усеченная модель, поэтому все, что нам нужно сделать, это квантовать ее при преобразовании в TFLite, как мы это делали раньше.
converter = tf.lite.TFLiteConverter.from_keras_model(model_pruned)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
model_pruned_quantized_tflite = converter.convert()
with open("model_pruned_quantized.tflite", "wb") as f:
f.write(model_pruned_quantized_tflite)
model_pruned_quantized_acc = evaluate_tflite_model(
"model_pruned_quantized.tflite", x_test, y_test
)
model_pruned_quantized_size = get_gzipped_model_size(
"model_pruned_quantized.tflite"
)
print(f"Pruned + Quantized accuracy: {model_pruned_quantized_acc}")
print(f"Pruned + Quantized size: {model_pruned_quantized_size}")
Pruned + Quantized accuracy: 0.9840168019364943
Pruned + Quantized size: 3996
Квантование не уменьшило точность усеченной модели, но уменьшило ее размер примерно на треть, с менее чем 6К до 4К байт.
Выводы
В целом, мы начали с базовой модели 14 КБ и в итоге получили сжатую версию 4 КБ, сокращение на 70%.
В то же время точность модели почти не ухудшилась.
Поскольку модели машинного обучения продолжают расти в размерах и развертываются на все более компактных устройствах, методы оптимизации моделей, такие как квантование и обрезка, становятся необходимыми для уменьшения размера модели и повышения ее эффективности. Реализуя эти методы в TensorFlow, мы можем сжимать модель без потери точности. Оптимизация моделей может привести к более быстрому выводу, меньшему объему памяти и повышению эффективности энергопотребления, что упрощает их развертывание на более широком спектре устройств.
Спасибо за прочтение!
Если вам понравился этот пост, почему бы вам не подписаться на обновления по электронной почте на мои новые статьи? А, став участником Medium, вы можете поддержать меня и получить неограниченный доступ ко всем историям других авторов и вашему покорному слуге.
Хотите всегда держать руку на пульсе стремительно развивающейся области машинного обучения и искусственного интеллекта? Ознакомьтесь с моим новым информационным бюллетенем AI Pulse. Нужна консультация? Вы можете спросить меня о чем угодно или заказать 1:1 здесь.
Вы также можете попробовать одну из других моих статей. Не можете выбрать? Выберите один из них:
Все изображения, если не указано иное, принадлежат автору.
