<jupyter_start><jupyter_text>Set default QuantDescriptor to use histogram based calibration for activation<jupyter_code>quant_desc_input = QuantDescriptor(calib_method='histogram')
quant_nn.QuantConv2d.set_default_quant_desc_input(quant_desc_input)
quant_nn.QuantLinear.set_default_quant_desc_input(quant_desc_input)<jupyter_output><empty_output><jupyter_text>Initialize quantized modules<jupyter_code>from pytorch_quantization import quant_modules
quant_modules.initialize()<jupyter_output><empty_output><jupyter_text>Create model with pretrained weight<jupyter_code>model = torchvision.models.resnet50(pretrained=True, progress=False)
model.cuda()<jupyter_output><empty_output><jupyter_text>Create data loader<jupyter_code>data_path = "/raid/data/imagenet/imagenet_pytorch"
batch_size = 512

traindir = os.path.join(data_path, 'train')
valdir = os.path.join(data_path, 'val')
_args = collections.namedtuple('mock_args', ['model', 'distributed', 'cache_dataset'])
dataset, dataset_test, train_sampler, test_sampler = load_data(traindir, valdir, _args(model='resnet50', distributed=False, cache_dataset=False))

data_loader = torch.utils.data.DataLoader(
    dataset, batch_size=batch_size,
    sampler=train_sampler, num_workers=4, pin_memory=True)

data_loader_test = torch.utils.data.DataLoader(
    dataset_test, batch_size=batch_size,
    sampler=test_sampler, num_workers=4, pin_memory=True)<jupyter_output>Loading data
Loading training data
Took 3.580507755279541
Loading validation data
Creating data loaders<jupyter_text>Calibrate the model<jupyter_code>def collect_stats(model, data_loader, num_batches):
    """Feed data to the network and collect statistic"""

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

    for i, (image, _) in tqdm(enumerate(data_loader), total=num_batches):
        model(image.cuda())
        if i >= num_batches:
            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 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()
# It is a bit slow since we collect histograms on CPU
with torch.no_grad():
    collect_stats(model, data_loader, num_batches=2)
    compute_amax(model, method="percentile", percentile=99.99)<jupyter_output>100%|██████████| 2/2 [04:50<00:00, 111.13s/it]<jupyter_text>Now evaluate the calibrated model<jupyter_code>criterion = nn.CrossEntropyLoss()
with torch.no_grad():
    evaluate(model, criterion, data_loader_test, device="cuda", print_freq=20)
    
# Save the model
torch.save(model.state_dict(), "/tmp/quant_resnet50-calibrated.pth")<jupyter_output>Test:  [ 0/98]  eta: 0:05:53  loss: 0.5656 (0.5656)  acc1: 85.7422 (85.7422)  acc5: 96.0938 (96.0938)  time: 3.6079  data: 2.8152  max mem: 5880
Test:  [20/98]  eta: 0:01:07  loss: 0.6741 (0.6825)  acc1: 82.8125 (82.4219)  acc5: 95.8984 (95.7682)  time: 0.7343  data: 0.0002  max mem: 5882
Test:  [40/98]  eta: 0:00:46  loss: 0.6995 (0.7157)  acc1: 80.0781 (81.4024)  acc5: 96.0938 (95.7412)  time: 0.7226  data: 0.0002  max mem: 5882
Test:  [60/98]  eta: 0:00:29  loss: 1.1064 (0.8590)  acc1: 71.4844 (78.2627)  acc5: 91.0156 (94.1150)  time: 0.7259  data: 0.0002  max mem: 5882
Test:  [80/98]  eta: 0:00:13  loss: 1.1220 (0.9372)  acc1: 72.4609 (76.7072)  acc5: 89.6484 (93.1375)  time: 0.7220  data: 0.0002  max mem: 5882
Test: Total time: 0:01:13
 * Acc@1 76.138 Acc@5 92.916<jupyter_text>We can also try different calibrations and see which one works the best<jupyter_code>with torch.no_grad():
    compute_amax(model, method="percentile", percentile=99.9)
    evaluate(model, criterion, data_loader_test, device="cuda", print_freq=20)
with torch.no_grad():
    for method in ["mse", "entropy"]:
        print(F"{method} calibration")
        compute_amax(model, method=method)
        evaluate(model, criterion, data_loader_test, device="cuda", print_freq=20)<jupyter_output>mse calibration
Test:  [ 0/98]  eta: 0:06:34  loss: 0.5700 (0.5700)  acc1: 85.1562 (85.1562)  acc5: 96.2891 (96.2891)  time: 4.0243  data: 3.3231  max mem: 5882
Test:  [20/98]  eta: 0:01:08  loss: 0.6758 (0.6838)  acc1: 82.8125 (82.5707)  acc5: 96.0938 (95.7868)  time: 0.7204  data: 0.0002  max mem: 5882
Test:  [40/98]  eta: 0:00:46  loss: 0.7047 (0.7163)  acc1: 80.2734 (81.4834)  acc5: 96.2891 (95.7746)  time: 0.7178  data: 0.0002  max mem: 5882
Test:  [60/98]  eta: 0:00:29  loss: 1.1127 (0.8585)  acc1: 71.0938 (78.3395)  acc5: 90.8203 (94.1278)  time: 0.7192  data: 0.0002  max mem: 5882
Test:  [80/98]  eta: 0:00:13  loss: 1.1261 (0.9367)  acc1: 72.6562 (76.7530)  acc5: 89.8438 (93.1785)  time: 0.7176  data: 0.0002  max mem: 5882
Test: Total time: 0:01:13
 * Acc@1 76.186 Acc@5 92.926
entropy calibration
Test:  [ 0/98]  eta: 0:05:28  loss: 0.5648 (0.5648)  acc1: 85.3516 (85.3516)  acc5: 96.0938 (96.0938)  time: 3.3558  data: 2.6268  max mem: 5882
Test:  [20/98]  eta: 0:01:05  loss: 0.6[...]
{'fraction_non_alphanumeric': 0.12976169768655418, 'fraction_numerical': 0.13097930074795616, 'mean_word_length': 3.8359966358284274, 'pattern_counts': {'":': 0, '<': 26, '<?xml version=': 0, '>': 26, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '7', 'dir': 'jupyter-structured-clean-dedup', 'id': '389524', 'license': 'Apache-2.0', 'n_tokens_mistral': 2668, 'n_tokens_neox': 2183, 'n_words': 487, 'path': 'tools/pytorch-quantization/examples/calibrate_quant_resnet50.ipynb', 'repo_name': 'malithj/TensorRT', 'size': '5749'}