【神经网络】如何在Pytorch中从零开始将MNIST网络量化为8位

论文：
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 = 1else:num_channels = 3self.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 = 64test_batch_size = 64epochs = 10lr = 0.01momentum = 0.5seed = 1log_interval = 500save_model = Falseno_cuda = Falseuse_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_intervalfor 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 = 0correct = 0with 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 losspred = output.argmax(dim=1, keepdim=True) # get the index of the max log-probabilitycorrect += 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 / scalezero_point = 0if initial_zero_point < qmin:zero_point = qminelif initial_zero_point > qmax:zero_point = qmaxelse:zero_point = initial_zero_pointzero_point = int(zero_point)q_x = zero_point + x / scaleq_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_nextzero_point_next = 0if initial_zero_point < qmin:zero_point_next = qminelif initial_zero_point > qmax:zero_point_next = qmaxelse:zero_point_next = initial_zero_pointzero_point_next = int(zero_point_next)return scale_next, zero_point_nextdef quantizeLayer(x, layer, stat, scale_x, zp_x):# for both conv and linear layersW = layer.weight.dataB = layer.bias.data# scale_x = x.scale# zp_x = x.zero_pointw = 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.scalezp_w = w.zero_pointscale_b = b.scalezp_b = b.zero_pointscale_next, zero_point_next = calcScaleZeroPoint(min_val=stat['min'], max_val=stat['max'])# Perparing input by shiftingX = x.float() - zp_xlayer.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 intx = layer(X) + zero_point_nextx = F.relu(x)# Resetlayer.weight.data = Wlayer.bias.data = Breturn x, scale_next, zero_point_nextdef quantForward(model, x, stats):# Quantise before inputting into incoming layersx = 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 layerx = 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笔记本，以运行和验证这个量化网络!

如果你们中有人想了解更多关于量化的知识，我已经把我从下面学到的资源嵌入其中。它们的确是无价的。