【神经网络】如何在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 = 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笔记本，以运行和验证这个量化网络!

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