Thanks for the reply.
Could you please let me know more detail about INT8 calibration?
or, any example code?
There’s no detail information regarding INT8 calibration on NVIDIA documentation.
It’s very difficult to refer these documents… only saying need calibration, only saying refer to documents…
where are those?
My code is as below. I got a help from Deploying Quantization Aware Trained models in INT8 using Torch-TensorRT — Torch-TensorRT v1.3.0 documentation. However, it doesn’t work. Please don’t say refer to forum…
import os
import torch
import torch_tensorrt
from torch_tensorrt import Input
from cc_weld_dataset import cc_weld_dataset
from unet_origin import UNet
from torch.utils.data import DataLoader
from tqdm import tqdm
from pytorch_quantization import nn as quant_nn
from pytorch_quantization import calib
def compute_amax(model, **kwargs):
# Load calib result
for name, module in model.named_modules():
if isinstance(module, quant_nn.TensorQuantizer):
if module._calibrator is not None:
if isinstance(module._calibrator, calib.MaxCalibrator):
module.load_calib_amax()
else:
module.load_calib_amax(**kwargs)
print(F"{name:40}: {module}")
model.cuda()
def collect_stats(model, data_loader, num_batches):
"""Feed data to the network and collect statistics"""
# Enable calibrators
for name, module in model.named_modules():
if isinstance(module, quant_nn.TensorQuantizer):
if module._calibrator is not None:
module.disable_quant()
module.enable_calib()
else:
module.disable()
# Feed data to the network for collecting stats
for i, (image, mask, original_image, image_path) in tqdm(enumerate(data_loader), total=num_batches):
# Forward pass with the image only
model(image.cuda())
if i >= num_batches - 1: # `i` is zero-based
break
# Disable calibrators
for name, module in model.named_modules():
if isinstance(module, quant_nn.TensorQuantizer):
if module._calibrator is not None:
module.enable_quant()
module.disable_calib()
else:
module.enable()
def calibrate_model(model, model_name, data_loader, num_calib_batch, calibrator, hist_percentile, out_dir):
"""
Feed data to the network and calibrate.
Arguments:
model: classification model
model_name: name to use when creating state files
data_loader: calibration data set
num_calib_batch: amount of calibration passes to perform
calibrator: type of calibration to use (max/histogram)
hist_percentile: percentiles to be used for historgram calibration
out_dir: dir to save state files in
"""
if num_calib_batch > 0:
print("Calibrating model")
with torch.no_grad():
collect_stats(model, data_loader, num_calib_batch)
if calibrator == "max":
compute_amax(model, method="max")
calib_output = os.path.join(
out_dir,
F"{model_name}-max-{num_calib_batch*data_loader.batch_size}.pth"
)
torch.save(model.state_dict(), calib_output)
elif calibrator == "histogram":
for percentile in hist_percentile:
print(F"{percentile} percentile calibration")
compute_amax(model, method="percentile", percentile=percentile)
calib_output = os.path.join(
out_dir,
F"{model_name}-percentile-{percentile}-{num_calib_batch*data_loader.batch_size}.pth"
)
torch.save(model.state_dict(), calib_output)
else:
raise ValueError(f"Unsupported calibrator type: {calibrator}")
print("Calibration completed!")
def convert_to_tensorrt_and_benchmark(model, input_shape=(8, 3, 256, 256), output_dir="./output"):
print("TensorRT starting...")
trt_compile_spec = {
"inputs": [torch_tensorrt.Input(input_shape)],
"enabled_precisions": {torch.int8},
}
trt_model = torch_tensorrt.compile(model, **trt_compile_spec)
print("TensorRT convert completed!")
trt_model_path = os.path.join(output_dir, "model_Unet_trt_int8.pt")
torch.jit.save(trt_model, trt_model_path)
print(f"TensorRT INT8 model save completed: {trt_model_path}")
print("TensorRT model inference benchmark starting...")
input_data = torch.randn(input_shape).to("cuda")
benchmark_model(trt_model, input_data)
def benchmark_model(model, input_data, num_warmup=10, num_runs=100):
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)
model.eval()
print("Warm-up...")
for _ in range(num_warmup):
with torch.no_grad():
_ = model(input_data)
print("inference time check starting...")
times = []
for _ in range(num_runs):
start_time = torch.cuda.Event(enable_timing=True)
end_time = torch.cuda.Event(enable_timing=True)
start_time.record()
with torch.no_grad():
_ = model(input_data)
end_time.record()
torch.cuda.synchronize()
times.append(start_time.elapsed_time(end_time))
print(f"average inference time (ms): {sum(times) / len(times):.2f}")
def main():
dataset_dir = "D:\\my_UNet"
output_dir = "D:\\my_UNet\\output"
os.makedirs(output_dir, exist_ok=True)
train_dataset = cc_weld_dataset(dataset_dir=os.path.join(dataset_dir, "dataset", "train"), augmentation=None)
train_loader = DataLoader(dataset=train_dataset, batch_size=8, shuffle=True)
model_path = os.path.join(output_dir, "model_Unet_origin_best.pt")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = UNet(n_channels=3, n_classes=1).to(device)
model.load_state_dict(torch.load(model_path, map_location=device))
model.eval()
print("Model Calibration starting...")
calibrate_model(
model=model,
model_name="unet",
data_loader=train_loader,
num_calib_batch=4,
calibrator="max",
hist_percentile=[99.9, 99.99, 99.999, 99.9999],
out_dir=output_dir
)
print("Model loading complted and TensorRT convert starting...")
convert_to_tensorrt_and_benchmark(
model=model,
input_shape=(8, 3, 256, 256),
output_dir=output_dir
)
if __name__ == "__main__":
main()