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

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也会训出比较大的模型，这个是不足的一个方面。

附录

1. fasttext language identification
2. fast源码
3. K-means学习总结
4. fasttext源码学习(1)–dictionary