7.2. Quick Start¶
The D-Robotics Plugin Pytorch (hereinafter referred to as Plugin) is based on the official PyTorch quantization interface and ideas. The Plugin adopts the Quantization Aware Training (QAT) solution, so users are recommended to read the PyTorch offical documentation related to the QAT.
The following figure shows the QAT overall process:
The following describes each step separately:
7.2.1. Building A Floating Point Model¶
When building a floating-point model, you need to pay attention to the following points:
Add quantization and dequantization nodes in the network. Generally, a quantization node needs to be inserted to the floating-point model at the beginning and a dequantization node needs to be inserted at the end. When a floating-point model has been converted to a QAT model that serves for Quantization Aware Training (QAT), the inserted quantization node will quantize the input.
Some floating-point operators of functional form need to be replaced by Pytorch’s operators inherited from the module or by some proprietary operators provided by Plugin.
Define the fusion function for floating-point operators and fuse the suitable operators.
There are more detailed requirements for building floating-point models in Tutorial.
import torch
import torch.optim as optim
import horizon_plugin_pytorch as D-Robotics
import os
from torch import nn
from torchvision import datasets, transforms
from torch.quantization import DeQuantStub
from horizon_plugin_pytorch.quantization import QuantStub, fuse_modules
class ConvBNReLU(nn.Sequential):
def __init__(self, in_channels, out_channels, kernel_size):
super(ConvBNReLU, self).__init__(
nn.Conv2d(
in_channels=in_channels,
out_channels=out_channels,
kernel_size=kernel_size
),
nn.BatchNorm2d(num_features=out_channels),
nn.ReLU()
)
def fuse_model(self):
fuse_modules(
self, ['0', '1', '2'], inplace=True,
)
class ClassiFier(nn.Module):
def __init__(self, in_channels, out_channels):
super(ClassiFier, self).__init__()
self.conv = nn.Conv2d(in_channels, out_channels, 1)
def forward(self, data):
return self.conv(data)
# Build a floating point model
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv0 = ConvBNReLU(1, 10, 5)
self.max_pool = nn.MaxPool2d(kernel_size=2)
self.conv1 = ConvBNReLU(10, 20, 5)
self.avg_pool = nn.AvgPool2d(kernel_size=8)
self.classifier = ClassiFier(20, 10)
# QuantStub's scale must be set to 1/128 explicitly to support BPU when getting input from the camera.
self.quant = QuantStub(scale=1/128)
self.dequant = DeQuantStub()
def forward(self, x):
# Insert a quantization node to quantize the input
x = self.quant(x)
x = self.conv0(x)
x = self.max_pool(x)
x = self.conv1(x)
x = self.avg_pool(x)
x = self.classifier(x)
# Insert a dequantization node to dequantize the output
x = self.dequant(x)
return x
# Define the fusion function
def fuse_model(self):
from horizon_plugin_pytorch import quantization
for m in self.modules():
if type(m) == ConvBNReLU:
m.fuse_model()
7.2.2. Pre-train A Floating Point Model¶
train_batch_size = 16
test_batch_size = 16
epoch_num = 1
neval_batches = 1
model_file = 'model.pt'
class AverageMeter(object):
"""Computes and stores the average and current value"""
def __init__(self, name, fmt=":f"):
self.name = name
self.fmt = fmt
self.reset()
def reset(self):
self.val = 0
self.avg = 0
self.sum = 0
self.count = 0
def update(self, val, n=1):
self.val = val
self.sum += val * n
self.count += n
self.avg = self.sum / self.count
def __str__(self):
fmtstr = "{name} {val" + self.fmt + "} ({avg" + self.fmt + "})"
return fmtstr.format(**self.__dict__)
criterion = nn.CrossEntropyLoss()
def accuracy(output, target, topk=(1,)):
"""Computes the accuracy over the k top predictions for the specified
values of k
"""
with torch.no_grad():
maxk = max(topk)
batch_size = target.size(0)
_, pred = output.topk(maxk, 1, True, True)
pred = pred.t()
correct = pred.eq(target.view(1, -1).expand_as(pred))
res = []
for k in topk:
correct_k = correct[:k].reshape(-1).float().sum(0, keepdim=True)
res.append(correct_k.mul_(100.0 / batch_size))
return res
def get_train_data_loader():
train_loader = torch.utils.data.DataLoader(
datasets.MNIST(
'mnist_data',
train=True,
download=True,
transform=transforms.Compose(
[transforms.ToTensor(),
transforms.Normalize((0.5,), (0.5,))]
)
),
batch_size=train_batch_size,
shuffle=True,
)
return train_loader
def get_test_data_loader():
train_loader = torch.utils.data.DataLoader(
datasets.MNIST(
'mnist_data',
train=False,
download=True,
transform=transforms.Compose(
[transforms.ToTensor(),
transforms.Normalize((0.5,), (0.5,))]
)
),
batch_size=test_batch_size,
shuffle=True,
)
return train_loader
data_loader = get_train_data_loader()
test_loader = get_test_data_loader()
def train(model, device, optimizer, epoch):
global min_loss
model.train()
for batch_idx, (data, target) in enumerate(data_loader):
data = data.to(device)
target = target.to(device)
output = model(data)
output = output.view(-1, 10)
loss = criterion(output, target)
optimizer.zero_grad()
loss.backward()
optimizer.step()
if batch_idx % 100 == 0:
print ('Train Epoch: {} batch {} \t Loss: {:.6f}'.
format(epoch, batch_idx, loss.item()))
def evaluate(model, device, neval_batches):
model.eval()
top1 = AverageMeter("Acc@1", ":6.2f")
top5 = AverageMeter("Acc@5", ":6.2f")
tested_batches = 0
with torch.no_grad():
for batch_idx, (data, target) in enumerate(test_loader):
tested_batches += 1
data = data.to(device)
target = target.to(device)
output = model(data)
output = output.view(-1, 10)
loss = criterion(output, target)
acc1, acc5 = accuracy(output, target, topk=(1, 5))
top1.update(acc1[0], data.size(0))
top5.update(acc5[0], data.size(0))
if tested_batches >= neval_batches:
return top1, top5
return top1, top5
def train_float_model(device):
model = Net().to(device)
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.1)
for nepoch in range(epoch_num):
train(model, device, optimizer, nepoch)
top1, top5 = evaluate(model, device, neval_batches)
print(
"float training Epoch %d :float evaluation accuracy on %d images, \
%2.2f" % (nepoch, neval_batches * test_batch_size, top1.avg)
)
torch.save(model.state_dict(), model_file)
train_float_model(torch.device('cuda'))
If you want to quantize the existing floating-point model, load the floating-point model first before operator fusion and quantization training. If the floating-point training is followed by the quantization aware training, there is no need to load it, and directly perform the subsequent steps.
def load_model():
model = Net()
state_dict = torch.load(model_file)
model.load_state_dict(state_dict)
model.to('cpu')
return model
qat_model = load_model()
7.2.3. Set BPU Architectures¶
horizon.march.set_march(horizon.march.March.XXX)
7.2.4. Operator Fusion¶
qat_model.fuse_model()
7.2.5. Convert A Floating Point Model To A Quantitative Model¶
def load_and_prepare_qat_model(device):
# Load a pre-trained floating point model
global qat_model
qat_model = qat_model.to(device)
top1, top5 = evaluate(qat_model, device, neval_batches)
print(
"float evaluation accuracy on %d images, \
%2.2f" % (neval_batches * test_batch_size, top1.avg)
)
# Set the quantization parameters for the QAT to specify how to quantize the weight and output of the operator.
qat_model.qconfig = horizon.quantization.get_default_qat_qconfig()
# Eliminate the quantization function of the output layer to improve the accuracy of the output.
qat_model.classifier.qconfig = \
horizon.quantization.get_default_qat_out_qconfig()
# Convert the floating point model to the quantized model.
horizon.quantization.prepare_qat(qat_model, inplace=True)
print(
"After preparation for QAT, note fake-quantization modules \n",
qat_model.conv0,
)
qat_model = qat_model.to(device)
load_and_prepare_qat_model(torch.device('cuda'))
7.2.6. Quantization Aware Training (QAT)¶
def quantization_training(device):
# Train a quantized model
optimizer = optim.SGD(qat_model.parameters(), lr=0.0001)
for nepoch in range(1):
train(qat_model, device, optimizer, nepoch)
# Train one epoch of quantized model for evaluation.
top1, top5 = evaluate(qat_model, device, neval_batches)
print(
"QAT Epoch %d :float evaluation accuracy on %d images, %2.2f"
% (nepoch, neval_batches * test_batch_size, top1.avg)
)
quantization_training(torch.device('cuda'))
7.2.7. Converting a Quantized Model to a Fixed-point Model¶
quantized_model = horizon.quantization.convert(
qat_model.eval(), inplace=False
)
7.2.8. Check and Compile a Fixed-point Prediction Model¶
def compile_quantized_model(device):
example_input = torch.ones(size=(neval_batches, 1, 28, 28), device=device)
traced_model = torch.jit.trace(quantized_model, example_input)
top1, top5 = evaluate(traced_model, device, neval_batches)
print(
"Traced : int evaluation accuracy on %d images, %2.2f"
% (neval_batches * test_batch_size, top1.avg)
)
# Check if the model can be compiled by hbdk. hbdk is a tool to compile fixed-point models.
horizon.quantization.check_model(quantized_model, example_input, advice=1)
hbdk_dir = "hbdk_model"
if not os.path.exists(hbdk_dir):
os.mkdir(hbdk_dir)
# Compile the model. The model.hbm in the hbdk_model directory is a compiled on-board model.
horizon.quantization.compile_model(
traced_model, [example_input], opt=2, hbm=hbdk_dir + "/model.hbm"
)
# Analyze the static performance on the model
horizon.quantization.perf_model(
traced_model,
[example_input],
opt=2,
input_source=["pyramid"],
layer_details=True,
out_dir=hbdk_dir,
)
horizon.quantization.visualize_model(
traced_model,
[example_input],
save_path=hbdk_dir + "/model.svg",
show=False,
)
compile_quantized_model(torch.device('cuda'))