fasttext源码学习(2)–模型压缩

前言

fasttext模型压缩的很明显，精度却降低不多，其网站上提供的语种识别模型，压缩前后的对比就是例证，压缩前126M，压缩后917K。太震惊了，必须学习一下。看文档介绍用到权重量化(weight quantization)和特征选择(feature selection)，下面结合代码学习下。

说明：文章中代码皆为简化版，为突出重点，简化了逻辑，原版代码需到官方网页下载。

一 特征选择

一开始以为fasttext会用到比较复杂的特征选择算法，直到看到代码才差点闪了腰。。。fasttext用的就是kbest，剩下的全砍掉，就是这么简单直接。

void FastText::quantize(const Args& qargs, const TrainCallback& callback) { if (qargs.cutoff > 0 && qargs.cutoff < input->size(0)) { auto idx = selectEmbeddings(qargs.cutoff); dict_->prune(idx); // 剪枝(词典重新计算) if (qargs.retrain) { // 重新训练 startThreads(callback); } } } std::vector<int32_t> FastText::selectEmbeddings(int32_t cutoff) const { std::shared_ptr<DenseMatrix> input = std::dynamic_pointer_cast<DenseMatrix>(input_); Vector norms(input->size(0)); input->l2NormRow(norms); // [1] 正则化 std::vector<int32_t> idx(input->size(0), 0); std::iota(idx.begin(), idx.end(), 0); std::sort(idx.begin(), idx.end(), [&norms, eosid](size_t i1, size_t i2) { return (eosid != i2 && norms[i1] > norms[i2]); }); // [2] 按正则化值排序 idx.erase(idx.begin() + cutoff, idx.end()); // [3] 保留指定数目的id return idx; }

从上面代码可以看出，selectEmbeddings主要干的就是按正则化值排序，然后只保留指定数目的行。默认保留是5万，结合上一章的DensMatrix的行数是200万+词语个数，可以看出这一步的特征选择最少能压缩到原来的1/40，压缩比很可观。

值得注意的是dict_->prune(idx)，这一步必不可少，因为fasttext是基于hash映射来计算矩阵下标的，特征被筛选后，相应的词典也需要清理压缩，对应关系也需要刷新。

void Dictionary::prune(std::vector<int32_t>& idx) { std::vector<int32_t> words, ngrams; // 按id选取对应的word、ngram for (auto it = idx.cbegin(); it != idx.cend(); ++it) { if (*it < nwords_) { words.push_back(*it); } else { ngrams.push_back(*it); } } // 按id排序 std::sort(words.begin(), words.end()); idx = words; // 计算被筛选的ngram的hash与id的对应关系 // ngram的对应关系原本不需存储，筛选后，由于对应关系的变化，导致需要存储pruneidx_ if (ngrams.size() != 0) { int32_t j = 0; for (const auto ngram : ngrams) { pruneidx_[ngram - nwords_] = j; j++; } idx.insert(idx.end(), ngrams.begin(), ngrams.end()); } pruneidx_size_ = pruneidx_.size(); int32_t j = 0; // 筛选过的word往前移，为后续的清除做准备 // 同时，重新计算hash与id的对应关系 for (int32_t i = 0; i < words_.size(); i++) { if (getType(i) == entry_type::label || (j < words.size() && words[j] == i)) { words_[j] = words_[i]; word2int_[find(words_[j].word)] = j; j++; } } nwords_ = words.size(); size_ = nwords_ + nlabels_; // 移除多余的word words_.erase(words_.begin() + size_, words_.end()); // 重新初始化各word的ngrams，重新计算ngram的下标 initNgrams(); }

要理解这段代码，必须先理解fasttext的数据存储即Dictionary和DenseMatrix那一部分，否则会非常晕。

二 权重量化

量化一般是将大的数值表示变为小的数值表示，比如从float变为byte。而fasttext采用了另外一种方法，product quantization。简单来说，就是将向量分割为更小的子向量，再使用kmeans算法，将子向量映射到中心点下标。这样, 假设子向量长度为2，则n*8的float类型的矩阵，被映射为n*4的byte矩阵，模型可以减小到原来的1/8.

void ProductQuantizer::train(int32_t n, const real* x) { std::vector<int32_t> perm(n, 0); std::iota(perm.begin(), perm.end(), 0); auto d = dsub_; auto np = std::min(n, max_points_); auto xslice = std::vector<real>(np * dsub_); // 划分为nsubq_个子向量，子向量长度为d. for (auto m = 0; m < nsubq_; m++) { std::shuffle(perm.begin(), perm.end(), rng); // 随机选取np行数据,每行长度为d for (auto j = 0; j < np; j++) { memcpy(xslice.data() + j * d, x + perm[j] * dim_ + m * dsub_, d*sizeof(real)); } // kmeans计算该子向量对应的中心点 kmeans(xslice.data(), get_centroids(m, 0), np, d); } } void ProductQuantizer::kmeans(const real* x, real* c, int32_t n, int32_t d) { std::vector<int32_t> perm(n, 0); std::iota(perm.begin(), perm.end(), 0); std::shuffle(perm.begin(), perm.end(), rng); // 随机初始化中心点 for (auto i = 0; i < ksub_; i++) { memcpy(&c[i * d], x + perm[i] * d, d * sizeof(real)); } auto codes = std::vector<uint8_t>(n); // kmeans标准算法，具体可参考另一篇介绍kmeans的文章 for (auto i = 0; i < niter_; i++) { Estep(x, c, codes.data(), d, n); MStep(x, c, codes.data(), d, n); } } // MStep中有一部分与Kmeans算法中不太一样的部分 // 对于中心点计数为0的部分做了修正，对中心点数值进行了调整 std::uniform_real_distribution<> runiform(0, 1); for (auto k = 0; k < ksub_; k++) { if (nelts[k] == 0) { int32_t m = 0; while (runiform(rng) * (n - ksub_) >= nelts[m] - 1) { m = (m + 1) % ksub_; } memcpy(centroids + k * d, centroids + m * d, sizeof(real) * d); for (auto j = 0; j < d; j++) { int32_t sign = (j % 2) * 2 - 1; centroids[k * d + j] += sign * eps_; centroids[m * d + j] -= sign * eps_; } nelts[k] = nelts[m] / 2; nelts[m] -= nelts[k]; } }

源码中默认子矩阵长度是2，kmeans簇大小为256(不超过一个byte)，默认会压缩至1/8大小。

总结

原本以为会写比较多内容，因为这部分代码确实花了点时间去看，尤其权重量化中kmeans算法前面那部分(子矩阵划分)，不太明白，但是看懂代码逻辑之后，又回头看了文档，恍然大悟，一句话就能把逻辑说的很清楚。但是并没办法绕过看代码，因为正是看文档看的不明白才去看代码的。。。-_-||

所以，最终整篇文章变成了代码注释，以防细节部分忘掉。。。

话说特征选择，个人觉得对于大部分场景，压缩比非常可观。但是也需要看到，DenseMatrix矩阵起步就是200万行，所以对于小数据集，fasttext也会训出比较大的模型，这个是不足的一个方面。

附录

fasttext language identificationfast源码K-means学习总结fasttext源码学习(1)–dictionary