Data-Free Quantization，我踩的那些坑
(本文首发于公众号，没事来逛逛)
上一篇文章介绍了高通 Data-Free Quantization 的基本思想，但我在代码实现的时候发现有个问题没有解决，因此这篇文章对前文打个补丁 (果然没有落实到代码层面都不能说自己看懂了论文)。
Depthwise Conv如何equalize
在前面的文章中，我给出了 weight equalize 的简单实现：
def equalize(weight1, bias1, weight2):
# 重排列
weight2 = weight2.permute(1, 0, 2, 3)
out_channel = weight1.shape[0]
for i in range(out_channel):
r1 = compute_range(weight1, i) # 计算kernel数值范围
r2 = compute_range(weight2, i)
s = r1 / sqrt(r1 * r2)
weight1[i] = weight1[i] * (1. / s)
weight2[i] = weight2[i] * s
bias1[i] = bias1[i] * (1. / s)
# 调整回之前的数据排布
weight2 = weight2.permute(1, 0, 2, 3)
return weight1, bias1, weight2
仔细读了代码的读者可能会发现：这段 equalize 的代码对可分离卷积 (depthwise conv) 好像不适用。
确实是这样的，上面这段代码只适用于普通的卷积。如果第二个 conv 是可分离卷积的话，由于它的 input channel 是 1，因此在循环里面，weight2 是会越界的。
高通在论文里面避开了这个问题，但其实可分离卷积才是需要 weight equalize 的地方。(太鸡贼了)
不过，好在高通自家的 AIMet 工具中已经落地了这套算法，所以在代码中可以了解到实际情况是怎么处理的。
首先，在 AIMet 中有这样一段代码 (https://github.com/quic/aimet/blob/develop/TrainingExtensions/torch/src/python/aimet_torch/cross_layer_equalization.py#L157 ，这个链接随着代码更新可能会变化)
@staticmethod
def convert_layer_group_to_cls_sets(layer_group):
"""
Helper function to convert a layer group to a list of cls sets
:param layer_group: Given layer group to conver
:return: List of cls sets
"""
cls_sets = []
prev_layer_to_scale = layer_group.pop(0)
while layer_group:
next_layer_to_scale = layer_group.pop(0)
if next_layer_to_scale.groups > 1:
if layer_group: # 如果第二个卷积是depthwise conv，则会继续找到下一个conv，三个为一组
next_non_depthwise_conv_layer = layer_group.pop(0)
cls_sets.append((prev_layer_to_scale, next_layer_to_scale, next_non_depthwise_conv_layer))
prev_layer_to_scale = next_non_depthwise_conv_layer
else:
cls_sets.append((prev_layer_to_scale, next_layer_to_scale))
prev_layer_to_scale = next_layer_to_scale
return cls_sets
这段代码会找到相邻的 conv，将他们组成一组后再做 weight equalize。如果是普通的卷积，则把相邻的两个 conv 作为一组进行 equalize，而如果第二个 conv 是 depthwise conv，则需要继续找到下一个相邻的卷积，然后三个一组 (conv, depthwise conv, conv)，一起 equalize。
从代码中变量名的命名，以及之后 equalize 的代码也可以看到，高通默认这里面的卷积不包括分组卷积 (所以要么 group = 1，要么就是彻底的 depthwise conv)，同时，高通也默认，如果遇到 depthwise conv，那紧跟在它后面的卷积一定是个普通的卷积 (即 group=1)，否则之后做 equalize 时就会报错。(私以为，这跟 weight equalize 在数学上的优化过程有关，如果没有满足这些条件，可能求解不出合适的缩放因子)。
找到这一组一组的卷积对后，就要开始做 weight equalize 了。这里也分两种情况，如果是两个普通的卷积为一组，那代码和前面的截图是一样的。而如果是三个卷积一组，根据官方给出的代码，可以总结为下面这张图：
跟之前 equalize 的区别是，这里有两个缩放因子 \(S^1\) 和 \(S^2\)。需要先用 \(S^1\) 对前两个 conv 进行 equalize，然后再用 \(S^2\) 对后两个 conv 进行 equalize。而且由于中间 depthwise conv 的特殊性，在 equalize 的时候不需要对它的 weight 进行重排 (关于重排，请参考上一篇文章)，只需要在后面的 equalize 中对第三个 conv 进行重排即可。
这两个缩放因子的计算方法和论文里介绍的没有本质区别，这里就不纠结数学上的细节了，直接给出答案：
其中，\(r_i^1\)、\(r_i^2\)、\(r_i^3\) 分别是三个卷积 weight 中第 \(i\) 个 kernel 的数值范围大小 (第三个卷积的 weight 需要重排)。
Talk is cheap，show you the code：
def equalize(weight1, bias1, weight2, bias2, weight3):
# 重排列
weight3 = weight3.permute(1, 0, 2, 3)
out_channel = weight1.shape[0]
S1, S2 = [], []
# 计算放缩系数
for i in range(out_channel):
r1 = compute_range(weight1, i) # 计算kernel数值范围
r2 = compute_range(weight2, i)
r3 = compute_range(weight3, i)
s = r1 / pow(r1 * r2 * r3, 1. / 3)
S1.append(s)
s = pow(r1 * r2 * r3, 1. / 3) / r3
S2.append(s)
# 对前两个conv进行equalize
for i in range(out_channel):
weight1[i] = weight1[i] * (1. / S1[i])
bias1[i] = bias1[i] * (1. / S1[i])
weight2[i] = weight2[i] * S1[i]
# 对后两个conv进行equalize
for i in range(out_channel):
weight2[i] = weight2[i] * (1. / S2[i])
bias2[i] = bias2[i] * (1. / S2[i])
weight3[i] = weight3[i] * S2[i]
# 调整回之前的数据排布
weight3 = weight3.permute(1, 0, 2, 3)
return weight1, bias1, weight2, bias2, weight3
顺便，我们看看高通官方 AIMet 的代码是怎么做的 (下一篇文章我会尝试用 pytorch fx 实现一下)。
在 AIMet 中，会先在 python 层面收集网络中的卷积，把它们组成一对一对，再通过 pybind 把卷积的 weight 和 bias 传递到 C++ 中进行 equalize。这篇文章主要看看 C++ 里面是如何实现我上面这段 equalize 代码的。
下面这段代码摘自 AIMet (https://github.com/quic/aimet/blob/develop/ModelOptimizations/DlEqualization/src/CrossLayerScaling.cpp#L93)
AimetEqualization::CrossLayerScaling::RescalingParamsVectors
CrossLayerScaling::scaleDepthWiseSeparableLayer(AimetEqualization::EqualizationParams& prevLayer, AimetEqualization::EqualizationParams& currLayer, AimetEqualization::EqualizationParams& nextLayer) {
const int ndims = 4;
int N = prevLayer.weightShape[0]; // output channels
// 获取三个卷积的weight和bias数据
cv::Mat weightTensor1 = cv::Mat(ndims, (int*) &prevLayer.weightShape[0], CV_32F, prevLayer.weight);
cv::Mat biasTensor1;
if (!prevLayer.isBiasNone)
biasTensor1 = cv::Mat(N, 1, CV_32F, (float*) &prevLayer.bias[0]);
cv::Mat weightTensor2 = cv::Mat(ndims, (int*) &currLayer.weightShape[0], CV_32F, currLayer.weight);
cv::Mat biasTensor2;
if (!currLayer.isBiasNone)
biasTensor2 = cv::Mat(N, 1, CV_32F, (float*) &currLayer.bias[0]);
cv::Mat weightTensor3 = cv::Mat(ndims, (int*) &nextLayer.weightShape[0], CV_32F, nextLayer.weight);
// 第三个卷积kernel重排
cv::Mat flippedWeightTensor3 = TensorOperations::swapFirstTwoAxisIn4dMat(weightTensor3);
// 计算缩放因子
RescalingParams* pReScalingMats = ScaleFactorCalculator::ForDepthWiseSeparableLayer(weightTensor1, weightTensor2, flippedWeightTensor3);
// 对前两个conv进行equalize
for (size_t s = 0; s < pReScalingMats->scalingMatrix12.total(); ++s) {
// Scaling Weight Matrix of prev layer with S12
cv::Mat w1PerChannel = TensorOperations::getDataPerChannelIn4dMat(weightTensor1, s, AXIS_0);
w1PerChannel = w1PerChannel * (1.0f / pReScalingMats->scalingMatrix12.at<float>(s));
// Scaling the bias of prev layer with S12
if (!prevLayer.isBiasNone)
biasTensor1.at<float>(s) = biasTensor1.at<float>(s) * (1.0f / pReScalingMats->scalingMatrix12.at<float>(s));
// Scaling Weight Matrix of curr layer with S12
cv::Mat w2PerChannel = TensorOperations::getDataPerChannelIn4dMat(weightTensor2, s, AXIS_0);
w2PerChannel = w2PerChannel * pReScalingMats->scalingMatrix12.at<float>(s);
}
// 对后两个conv进行equalize
for (size_t s = 0; s < pReScalingMats->scalingMatrix23.total(); ++s) {
// Scaling Weight Matrix of prev layer with S23
cv::Mat w2PerChannel = TensorOperations::getDataPerChannelIn4dMat(weightTensor2, s, AXIS_0);
w2PerChannel = w2PerChannel * (1.0f / pReScalingMats->scalingMatrix23.at<float>(s));
// Scaling the bias of curr layer with S23
if (!currLayer.isBiasNone)
biasTensor2.at<float>(s) = biasTensor2.at<float>(s) * (1.0f / pReScalingMats->scalingMatrix23.at<float>(s));
// Scaling Weight Matrix of curr layer with S23
cv::Mat w3PerChannel = TensorOperations::getDataPerChannelIn4dMat(flippedWeightTensor3, s, AXIS_0);
w3PerChannel = w3PerChannel * pReScalingMats->scalingMatrix23.at<float>(s);
}
// 调整回之前的数据排布
cv::Mat(TensorOperations::swapFirstTwoAxisIn4dMat(flippedWeightTensor3)).copyTo(weightTensor3);
// return pReScalingMats as vectors
CrossLayerScaling::RescalingParamsVectors scalingVectors;
scalingVectors.scalingMatrix12.assign(pReScalingMats->scalingMatrix12.begin<float>(),
pReScalingMats->scalingMatrix12.end<float>());
scalingVectors.scalingMatrix23.assign(pReScalingMats->scalingMatrix23.begin<float>(),
pReScalingMats->scalingMatrix23.end<float>());
return scalingVectors;
}
上面关键步骤我加了注释，可以看到和我前面给出的代码基本是一样的流程。
然后是具体计算缩放因子的代码 (https://github.com/quic/aimet/blob/develop/ModelOptimizations/DlEqualization/src/ScaleFactorCalculator.cpp#L90)
AimetEqualization::RescalingParams* ScaleFactorCalculator::ForDepthWiseSeparableLayer(const cv::Mat& weightTensor1, const cv::Mat& weightTensor2, const cv::Mat& weightTensor3)
{
AimetEqualization::RescalingParams* reScalingMats = new RescalingParams;
// 省略若干代码....
// 分别计算三组weight里面每个kernel的数值范围
cv::Mat rangeVec1 = TensorOperations::computeRangeAlongFirstAxis(weightTensor1);
cv::Mat rangeVec2 = TensorOperations::computeRangeAlongFirstAxis(weightTensor2);
cv::Mat rangeVec3 = TensorOperations::computeRangeAlongFirstAxis(weightTensor3);
// 三次开方计算缩放系数
cv::Mat cubeRootMat;
// perform element-wise multiplication on range vectors and find sqrt
cv::pow((rangeVec1.mul(rangeVec2).mul(rangeVec3)), 1.0f / 3, cubeRootMat);
reScalingMats->scalingMatrix12 = cv::Mat::ones(1, rangeVec1.total(), FLOAT_32_TYPE);
reScalingMats->scalingMatrix23 = cv::Mat::ones(1, rangeVec2.total(), FLOAT_32_TYPE);
// 计算第一个缩放因子
for (size_t s = 0; s < rangeVec1.total(); ++s) {
if ((rangeVec1.at<float>(s) != 0) && (rangeVec2.at<float>(s) != 0) && (rangeVec3.at<float>(s) != 0)) {
reScalingMats->scalingMatrix12.at<float>(s) = (rangeVec1.at<float>(s)) * (1.0f / cubeRootMat.at<float>(s));
}
}
// 计算第二个缩放因子
for (size_t s = 0; s < rangeVec2.total(); ++s) {
if ((rangeVec1.at<float>(s) != 0) && (rangeVec2.at<float>(s) != 0) && (rangeVec3.at<float>(s) != 0)) {
reScalingMats->scalingMatrix23.at<float>(s) = (cubeRootMat.at<float>(s)) * (1.0f / rangeVec3.at<float>(s));
}
}
return reScalingMats;
}
总结
这篇文章介绍了 Data-Free quantization 如何对 depthwise conv 进行 weight equalize，算是对前文的补充介绍，并简单介绍了官方的代码实现。下篇文章中我会自己用 pytorch 实现一遍 weight equalize，顺便也看看 bias correction 的实现。
参考
欢迎关注我的公众号：大白话AI，立志用大白话讲懂AI。
