<jupyter_start><jupyter_text>This notebook contains an excerpt from the [PyTorch Pocket Reference](http://pytorchbook.com) book by [<NAME>](http://joepapa.ai); content is available [on GitHub](https://github.com/joe-papa/pytorch-book). [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/joe-papa/pytorch-book/blob/main/06_04_Quantization.ipynb)  Chapter 6 - PyTorch Acceleration & Optimization We use CPU backend for this example  Quantization<jupyter_code>import torch 
from torch import nn
import torch.nn.functional as F

class LeNet5(nn.Module):
    def __init__(self):
        super(LeNet5, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = F.max_pool2d(
            F.relu(self.conv1(x)), (2, 2))
        x = F.max_pool2d(
            F.relu(self.conv2(x)), 2)
        x = x.view(-1, 
                   int(x.nelement() / x.shape[0]))
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

model = LeNet5()
for n, p in model.named_parameters():
  print(n, ": ", p.dtype)

# out:
# conv1.weight :  torch.float32
# conv1.bias :  torch.float32
# conv2.weight :  torch.float32
# conv2.bias :  torch.float32
# fc1.weight :  torch.float32
# fc1.bias :  torch.float32
# fc2.weight :  torch.float32
# fc2.bias :  torch.float32
# fc3.weight :  torch.float32
# fc3.bias :  torch.float32
model = model.half()

for n, p in model.named_parameters():
  print(n, ": ", p.dtype)

# out: 
# conv1.weight :  torch.float16
# conv1.bias :  torch.float16
# conv2.weight :  torch.float16
# conv2.bias :  torch.float16
# fc1.weight :  torch.float16
# fc1.bias :  torch.float16
# fc2.weight :  torch.float16
# fc2.bias :  torch.float16
# fc3.weight :  torch.float16
# fc3.bias :  torch.float16
import torch.quantization

quantized_model = \
  torch.quantization.quantize_dynamic(
      model, 
      {torch.nn.Linear},
      dtype=torch.qint8)
static_quant_model = LeNet5()
static_quant_model.qconfig = \
  torch.quantization.get_default_qconfig('fbgemm')

torch.quantization.prepare(
    static_quant_model, inplace=True)
torch.quantization.convert(
    static_quant_model, inplace=True)
qat_model = LeNet5()
qat_mode.qconfig = \
  torch.quantization.get_default_qat_qconfig('fbgemm')

torch.quantization.prepare_qat(
    qat_model, inplace=True)
torch.quantization.convert(
    qat_model, inplace=True)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.13961407491486946, 'fraction_numerical': 0.04275444570563754, 'mean_word_length': 3.257648953301127, 'pattern_counts': {'":': 2, '<': 6, '<?xml version=': 0, '>': 6, 'https://': 3, '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': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '664653', 'license': 'MIT', 'n_tokens_mistral': 1064, 'n_tokens_neox': 1044, 'n_words': 179, 'path': '06_04_Quantization.ipynb', 'repo_name': 'kingslyt/pytorch-book', 'size': '2645'}