Pytorch quantitative perception training (qat) steps

# QAT follows the same steps as PTQ, with the exception of the training loop before you actually convert the model to its quantized version
# QAT Follow and PTQ The same procedure , In addition to the training cycle before actually converting the model to a quantitative version 
''''''
''' Quantify perception training steps ： step1. Build a model  step2. The fusion （ An optional step ） step3. Insert stubs（1 and 3 Can be combined ） step4. Get ready （ The main thing is to choose the architecture ） step5. Training  step6. Model transformation  '''
import torch
from torch import nn

backend = "fbgemm"  # running on a x86 CPU. Use "qnnpack" if running on ARM.

'''step1. Build a model build model'''
m = nn.Sequential(
     nn.Conv2d(2,64,8),
     nn.ReLU(),
     nn.Conv2d(64, 128, 8),
     nn.ReLU(),
)

"""step2. The fusion Fuse（ An optional step ）"""
torch.quantization.fuse_modules(m, ['0','1'], inplace=True) # fuse first Conv-ReLU pair
torch.quantization.fuse_modules(m, ['2','3'], inplace=True) # fuse second Conv-ReLU pair

"""step3. Insert stubs In the model ,Insert stubs"""
m = nn.Sequential(torch.quantization.QuantStub(),
                  *m,
                  torch.quantization.DeQuantStub())

"""step4. Get ready Prepare"""
m.train()
m.qconfig = torch.quantization.get_default_qconfig(backend)
torch.quantization.prepare_qat(m, inplace=True)

"""step5. Training Training Loop"""
n_epochs = 10
opt = torch.optim.SGD(m.parameters(), lr=0.1)
loss_fn = lambda out, tgt: torch.pow(tgt-out, 2).mean()
for epoch in range(n_epochs):
  x = torch.rand(10,2,24,24)
  out = m(x)
  loss = loss_fn(out, torch.rand_like(out))
  opt.zero_grad()
  loss.backward()
  opt.step()
  print(loss)

"""step6. Model transformation Convert"""
m.eval()
torch.quantization.convert(m, inplace=True)