Introduction
Dynamic quantization quantize the weights of neural networks to integers, but the activations are dynamically quantized during inference. Comparing to floating point neural networks, the size of dynamic quantized model is much smaller since the weights are stored as low-bitwidth integers. Comparing to other quantization techniques, dynamic quantization does not require any data for calibration or fine-tuning. More details about the mathematical foundations of quantization for neural networks could be found in my article “Quantization for Neural Networks”.
Given a pre-trained floating point model, we could easily create an dynamically quantized model, run inference, and potentially achieve better latency without too much additional effort. In this blog post, I would like to show how to use PyTorch to do dynamic quantizations.
PyTorch Dynamic Quantization
Unlike TensorFlow 2.3.0 which supports integer quantization using arbitrary bitwidth from 2 to 16, PyTorch 1.7.0 only supports 8-bit integer quantization. The workflow is as easy as loading a pre-trained floating point model and apply a dynamic quantization wrapper.
In this case, I would like to use the BERT-QA model from HuggingFace Transformers as an example. I was dynamically quantizing the torch.nn.Linear layer for the BERT-QA model since the majority of the computation for Transformer based models are matrix multiplications. The source code could also be downloaded from GitHub.
# qa.py
import os
import time
import torch
from transformers import BertTokenizer, BertForQuestionAnswering
def measure_inference_latency(model, inputs, num_samples=100, num_warmups=10):
with torch.no_grad():
for _ in range(num_warmups):
_ = model(**inputs)
torch.cuda.synchronize()
with torch.no_grad():
start_time = time.time()
for _ in range(num_samples):
_ = model(**inputs)
torch.cuda.synchronize()
end_time = time.time()
elapsed_time = end_time - start_time
elapsed_time_ave = elapsed_time / num_samples
return elapsed_time_ave
def get_bert_qa_model(model_name="deepset/bert-base-cased-squad2", cache_dir="./saved_models"):
# https://huggingface.co/transformers/model_doc/bert.html#transformers.BertForQuestionAnswering
tokenizer = BertTokenizer.from_pretrained(model_name, cache_dir=cache_dir)
model = BertForQuestionAnswering.from_pretrained(model_name, cache_dir=cache_dir, return_dict=True)
return model, tokenizer
def prepare_qa_inputs(question, text, tokenizer, device=None):
inputs = tokenizer(question, text, return_tensors="pt")
if device is not None:
inputs_cuda = dict()
for input_name in inputs.keys():
inputs_cuda[input_name] = inputs[input_name].to(device)
inputs = inputs_cuda
return inputs
def move_inputs_to_device(inputs, device=None):
inputs_cuda = dict()
for input_name in inputs.keys():
inputs_cuda[input_name] = inputs[input_name].to(device)
return inputs_cuda
def run_qa(model, tokenizer, question, text, device=None):
inputs = prepare_qa_inputs(question=question, text=text, tokenizer=tokenizer)
all_tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"].numpy()[0])
if device is not None:
inputs = move_inputs_to_device(inputs, device=device)
model = model.to(device)
outputs = model(**inputs)
start_scores = outputs.start_logits
end_scores = outputs.end_logits
answer_start_idx = torch.argmax(start_scores, 1)[0]
answer_end_idx = torch.argmax(end_scores, 1)[0] + 1
answer = " ".join(all_tokens[answer_start_idx : answer_end_idx])
return answer
def get_model_size(model, temp_dir="/tmp"):
model_dir = os.path.join(temp_dir, "temp")
torch.save(model.state_dict(), model_dir)
# model.save_pretrained(model_dir)
size = os.path.getsize(model_dir)
os.remove(model_dir)
return size
def main():
cuda_device = torch.device("cuda:0")
num_samples = 100
model, tokenizer = get_bert_qa_model(model_name="deepset/bert-base-cased-squad2")
model.eval()
# https://pytorch.org/docs/stable/torch.quantization.html?highlight=torch%20quantization%20quantize_dynamic#torch.quantization.quantize_dynamic
quantized_model = torch.quantization.quantize_dynamic(model, {torch.nn.Linear}, dtype=torch.qint8)
print("=" * 75)
print("Model Sizes")
print("=" * 75)
model_size = get_model_size(model=model)
quantized_model_size = get_model_size(model=quantized_model)
print("FP32 Model Size: {:.2f} MB".format(model_size / (2 ** 20)))
print("INT8 Model Size: {:.2f} MB".format(quantized_model_size / (2 ** 20)))
question = "What publication printed that the wealthiest 1% have more money than those in the bottom 90%?"
text = "According to PolitiFact the top 400 richest Americans \"have more wealth than half of all Americans combined.\" According to the New York Times on July 22, 2014, the \"richest 1 percent in the United States now own more wealth than the bottom 90 percent\". Inherited wealth may help explain why many Americans who have become rich may have had a \"substantial head start\". In September 2012, according to the Institute for Policy Studies, \"over 60 percent\" of the Forbes richest 400 Americans \"grew up in substantial privilege\"."
inputs = prepare_qa_inputs(question=question, text=text, tokenizer=tokenizer)
answer = run_qa(model=model, tokenizer=tokenizer, question=question, text=text)
answer_quantized = run_qa(model=quantized_model, tokenizer=tokenizer, question=question, text=text)
print("=" * 75)
print("BERT QA Example")
print("=" * 75)
print("Text: ")
print(text)
print("Question: ")
print(question)
print("Model Answer: ")
print(answer)
print("Dynamic Quantized Model Answer: ")
print(answer_quantized)
print("=" * 75)
print("BERT QA Inference Latencies")
print("=" * 75)
model_latency = measure_inference_latency(model=model, inputs=inputs, num_samples=num_samples)
print("CPU Inference Latency: {:.2f} ms / sample".format(model_latency * 1000))
quantized_model_latency = measure_inference_latency(model=quantized_model, inputs=inputs, num_samples=num_samples)
print("Dynamic Quantized CPU Inference Latency: {:.2f} ms / sample".format(quantized_model_latency * 1000))
inputs_cuda = move_inputs_to_device(inputs, device=cuda_device)
model.to(cuda_device)
model_cuda_latency = measure_inference_latency(model=model, inputs=inputs_cuda, num_samples=num_samples)
print("CUDA Inference Latency: {:.2f} ms / sample".format(model_cuda_latency * 1000))
# No CUDA backend for dynamic quantization in PyTorch 1.7.0
# quantized_model_cuda = quantized_model.to(cuda_device)
# quantized_model_cuda_latency = measure_inference_latency(model=quantized_model_cuda, inputs=inputs_cuda, num_samples=num_samples)
# print("Dynamic Quantized GPU Inference Latency: {:.2f} ms / sample".format(quantized_model_cuda_latency * 1000))
if __name__ == "__main__":
main()
With PyTorch 1.7.0, we could do dynamic quantization using x86-64 and aarch64 CPUs. However, NVIDIA GPUs have not been supported for PyTorch dynamic quantization yet.
$ python qa.py
===========================================================================
Model Sizes
===========================================================================
FP32 Model Size: 411.00 MB
INT8 Model Size: 168.05 MB
===========================================================================
BERT QA Example
===========================================================================
Text:
According to PolitiFact the top 400 richest Americans "have more wealth than half of all Americans combined." According to the New York Times on July 22, 2014, the "richest 1 percent in the United States now own more wealth than the bottom 90 percent". Inherited wealth may help explain why many Americans who have become rich may have had a "substantial head start". In September 2012, according to the Institute for Policy Studies, "over 60 percent" of the Forbes richest 400 Americans "grew up in substantial privilege".
Question:
What publication printed that the wealthiest 1% have more money than those in the bottom 90%?
Model Answer:
New York Times
Dynamic Quantized Model Answer:
New York Times
===========================================================================
BERT QA Inference Latencies
===========================================================================
CPU Inference Latency: 52.27 ms / sample
Dynamic Quantized CPU Inference Latency: 40.63 ms / sample
CUDA Inference Latency: 7.02 ms / sample
We could see that the model size of the INT8 quantized model is much smaller than the FP32 model. The inference latency of INT8 dynamic quantization on CPU is much faster than the FP32 ordinary inference on CPU. However, FP32 inference using NVIDIA GPU is still the fastest.