论文：
Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only Inference
下载地址：https://arxiv.org/pdf/1712.05877.pdf
更新:量化感知训练的博客文章是在线的，并在这里链接，通过它我们可以训练和量化我们的模型以运行在4比特!
你好，我想分享我如何能够使用定点算术(8位算术)运行神经网络推理的旅程。截至目前，Pytorch的状态只允许32位或16位浮点训练和推理。如果现在想要使用量化压缩Pytorch神经网络，他/她需要将其导入onnx，转换为caffe，并在计算图上运行辉光量化编译器，最终产生量化网络。
在深入研究如何量化一个网络之前，让我们看看为什么我们需要量化一个网络。简单的答案是提高推理速度，浮点运算通常比定点(整数)运算需要更长的计算时间。另一个优势是节省空间，浮点网络的大小是8位量化网络的4倍。这与边缘设备(手机、物联网)尤其相关，因为低存储空间和计算需求对其成为可生产的解决方案至关重要。
在继续之前，这里有一个工作的Colab笔记本，供那些只想查看代码的人运行和验证这个量化网络。此示例在普通Pytorch中从头开始实现量化(没有外部库或框架)
现在我们已经证明了量化的必要性，让我们看看如何量化一个简单的MNIST模型。让我们使用一个简单的模型架构来解决MNIST，它使用2个conv层和2个全连接层。
class Net(nn.Module):
def __init__(self, mnist=True):
super(Net, self).__init__()
if mnist:
num_channels = 1
else:
num_channels = 3
self.conv1 = nn.Conv2d(num_channels, 20, 5, 1)
self.conv2 = nn.Conv2d(20, 50, 5, 1)
self.fc1 = nn.Linear(4*4*50, 500)
self.fc2 = nn.Linear(500, 10)
def forward(self, x):
x = F.relu(self.conv1(x))
x = F.max_pool2d(x, 2, 2)
x = F.relu(self.conv2(x))
x = F.max_pool2d(x, 2, 2)
x = x.view(-1, 4*4*50)
x = F.relu(self.fc1(x))
x = self.fc2(x)
return F.log_softmax(x, dim=1)
让我们使用下面的简单训练脚本来训练这个网络:
def train(args, model, device, train_loader, optimizer, epoch):
model.train()
for batch_idx, (data, target) in enumerate(train_loader):
data, target = data.to(device), target.to(device)
optimizer.zero_grad()
output = model(data)
loss = F.nll_loss(output, target)
loss.backward()
optimizer.step()
if batch_idx % args["log_interval"] == 0:
print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
epoch, batch_idx * len(data), len(train_loader.dataset),
100. * batch_idx / len(train_loader), loss.item()))
def main():
batch_size = 64
test_batch_size = 64
epochs = 10
lr = 0.01
momentum = 0.5
seed = 1
log_interval = 500
save_model = False
no_cuda = False
use_cuda = not no_cuda and torch.cuda.is_available()
torch.manual_seed(seed)
device = torch.device("cuda" if use_cuda else "cpu")
kwargs = {'num_workers': 1, 'pin_memory': True} if use_cuda else {}
train_loader = torch.utils.data.DataLoader(
datasets.MNIST('../data', train=True, download=True,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])),
batch_size=batch_size, shuffle=True, **kwargs)
test_loader = torch.utils.data.DataLoader(
datasets.MNIST('../data', train=False, transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])),
batch_size=test_batch_size, shuffle=True, **kwargs)
model = Net().to(device)
optimizer = optim.SGD(model.parameters(), lr=lr, momentum=momentum)
args = {}
args["log_interval"] = log_interval
for epoch in range(1, epochs + 1):
train(args, model, device, train_loader, optimizer, epoch)
test(args, model, device, test_loader)
if (save_model):
torch.save(model.state_dict(),"mnist_cnn.pt")
return model
现在，我们可以使用简单的```model = main()``命令来训练这个网络。一旦模型被训练了10个epoch，让我们通过以下测试函数来测试这个模型。
def test(args, model, device, test_loader):
model.eval()
test_loss = 0
correct = 0
with torch.no_grad():
for data, target in test_loader:
data, target = data.to(device), target.to(device)
output = model(data)
test_loss += F.nll_loss(output, target, reduction='sum').item() # sum up batch loss
pred = output.argmax(dim=1, keepdim=True) # get the index of the max log-probability
correct += pred.eq(target.view_as(pred)).sum().item()
test_loss /= len(test_loader.dataset)
print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
test_loss, correct, len(test_loader.dataset),
100. * correct / len(test_loader.dataset)))
经过测试，我们得到了99%的准确率。(~9900/10000)正确分类。现在让我们研究一下通过一种称为后训练量化的技术来执行量化。
要点是我们将神经网络的激活和权重转换为8位整数(范围为0到255)。因此，我们在不动点上执行所有算法，希望精度不会显著下降。
为了量化和去量化一个张量，我们使用以下公式:
x_Float =缩放*(x_Quant -zero_point)。因此,
x_Quant = (x_Float/scale) +零点。
这里缩放等于(max_val - min_val) / (qmax - qmin)
其中max_val和min_val分别是X张量的最大值和最小值。Qmin和q_max表示8位数字的范围(分别为0和255)。刻度会缩放量化网络，零点会移动数字。下面给出的去量化和量化函数更清楚地说明了浮点张量如何转换为8位张量，反之亦然。
QTensor = namedtuple('QTensor', ['tensor', 'scale', 'zero_point'])
def quantize_tensor(x, num_bits=8):
qmin = 0.
qmax = 2.**num_bits - 1.
min_val, max_val = x.min(), x.max()
scale = (max_val - min_val) / (qmax - qmin)
initial_zero_point = qmin - min_val / scale
zero_point = 0
if initial_zero_point < qmin:
zero_point = qmin
elif initial_zero_point > qmax:
zero_point = qmax
else:
zero_point = initial_zero_point
zero_point = int(zero_point)
q_x = zero_point + x / scale
q_x.clamp_(qmin, qmax).round_()
q_x = q_x.round().byte()
return QTensor(tensor=q_x, scale=scale, zero_point=zero_point)
def dequantize_tensor(q_x):
return q_x.scale * (q_x.tensor.float() - q_x.zero_point)
需要注意的是，刻度是浮点数，而零点是整数(8位)。然而，现代实现通过一些花哨的位技巧(即近似)绕过了这种规模的浮点乘法，这些技巧被证明对网络的精度影响可以忽略不计。
现在我们已经准备好了这些函数，我们可以通过修改MNIST网络的前向传递来量化我们的权重和激活。修改后的前向传球看起来像这样。
def calcScaleZeroPoint(min_val, max_val,num_bits=8):
# Calc Scale and zero point of next
qmin = 0.
qmax = 2.**num_bits - 1.
scale_next = (max_val - min_val) / (qmax - qmin)
initial_zero_point = qmin - min_val / scale_next
zero_point_next = 0
if initial_zero_point < qmin:
zero_point_next = qmin
elif initial_zero_point > qmax:
zero_point_next = qmax
else:
zero_point_next = initial_zero_point
zero_point_next = int(zero_point_next)
return scale_next, zero_point_next
def quantizeLayer(x, layer, stat, scale_x, zp_x):
# for both conv and linear layers
W = layer.weight.data
B = layer.bias.data
# scale_x = x.scale
# zp_x = x.zero_point
w = quantize_tensor(layer.weight.data)
b = quantize_tensor(layer.bias.data)
layer.weight.data = w.tensor.float()
layer.bias.data = b.tensor.float()
####################################################################
# This is Quantisation !!!!!!!!!!!!!!!!!!!!!!!!!!!!!
scale_w = w.scale
zp_w = w.zero_point
scale_b = b.scale
zp_b = b.zero_point
scale_next, zero_point_next = calcScaleZeroPoint(min_val=stat['min'], max_val=stat['max'])
# Perparing input by shifting
X = x.float() - zp_x
layer.weight.data = (scale_x * scale_w/scale_next)*(layer.weight.data - zp_w)
layer.bias.data = (scale_b/scale_next)*(layer.bias.data + zp_b)
# All int
x = layer(X) + zero_point_next
x = F.relu(x)
# Reset
layer.weight.data = W
layer.bias.data = B
return x, scale_next, zero_point_next
def quantForward(model, x, stats):
# Quantise before inputting into incoming layers
x = quantize_tensor_act(x, stats['conv1'])
x, scale_next, zero_point_next = quantizeLayer(x.tensor, model.conv1, stats['conv2'], x.scale, x.zero_point)
x = F.max_pool2d(x, 2, 2)
x, scale_next, zero_point_next = quantizeLayer(x, model.conv2, stats['fc1'], scale_next, zero_point_next)
x = F.max_pool2d(x, 2, 2)
x = x.view(-1, 4*4*50)
x, scale_next, zero_point_next = quantizeLayer(x, model.fc1, stats['fc2'], scale_next, zero_point_next)
# Back to dequant for final layer
x = dequantize_tensor(QTensor(tensor=x, scale=scale_next, zero_point=zero_point_next))
x = model.fc2(x)
return F.log_softmax(x, dim=1)
在这里，我们在输入卷积层conv1之前对激活进行量化，并使用名为quantizeLayer的函数，该函数接受conv或线性层以及量化激活的激活、缩放和零点，quantizeLayer()函数执行完全量化的层的前向传递。如果您有任何疑问，请查看上面的代码。您可能想知道quantize_tensor_act()函数是做什么的，它只是通过遍历1000个示例并平均结果，使用张量x通常具有的最小值和最大值对激活x进行量化。它使用这些统计数据来计算尺度，从而计算零点，这是量化张量的必要条件。现在，让我们将所有这些放在一起，并使用这种新的quantForward方法运行网络，并检查最终的准确性。
仍然是99% !当然，这只是一个玩具例子，我已经严重跳过了量化理论，但这是神经网络中如何执行量化的基本要点。它不是巫毒魔法，而是简单的线性代数和一些巧妙的技巧来绕过pytorch层。
希望这对你们来说是一个有趣的旅程，请查看这个正在工作的Colab笔记本，以运行和验证这个量化网络!
如果你们中有人想了解更多关于量化的知识，我已经把我从下面学到的资源嵌入其中。它们的确是无价的。