Ternary weight networks
Introduction
这两天看了一下这篇文章,我就这里分享一下,不过我还是只记录一下跟别人blog上没有,或者自己的想法(ps: 因为有时候翻blog时候发现每篇都一样还是挺烦的= =) 。为了不重复前人的工作,我post一个不小心翻到的博客权值简化(1):三值神经网络(Ternary Weight Networks),整个论文内容及实现都讲的很全面了,可以翻阅一下,我也借鉴一下。
文中主要工作的点在三个方面:
- 增加了网络的表达力(expressive ability)。在{1,0,1}基础上增加了 \(\alpha\) 作为scaled factor;
- 压缩模型大小。当然主要是weight的压缩。比起FPWN(full precision weight network)有16~32x的提升,但是BPWN(binary precision weight network)的2x大小(ps:当然在TWN的caffe代码里面,都由float double类型存储,因为这需要在应该上方面来实现);
- 减少计算需求。主要相比于BPWN增多了0,当然这方面也需硬件来获得提升,在该caffe代码里面并没有;
Ternary Quantization
在我的理解看来,文中最核心的内容是:将有约束的并且两变量之间互相依赖的优化问题,逐步拆分最后用具有先验的统计方法来近视解决。
最初的优化问题:
将\(W^{t}\)的约束具体化为:
并将其带入公式(1),将\(W^{t*}\)的优化转化为\(\Delta^*\)的优化:
然后对公式(4)中的\(\alpha\)求偏导,得到:
因为\(\alpha\)和\(\Delta\)相互依赖,将(5)代入(4)消去\(\alpha\):
但问题来了,公式(6)依然没法求,而文中就根据先验知识,假设\(W_i\)服从\(N(0,\sigma^2)\)分布,近视的\(\Delta^*\)为\(0.6\sigma\)(\(0.6\sigma\)等于\(0.75E(|W|)\))。因此作者采用粗暴的方法,把\(\Delta^*\)设为\(\Delta^*\approx0.7E(|W|)\approx\frac{n}{0.7}\sum_{i=1}^n|W_i|\)
//caffe-twns
//blob.cpp
template <typename Dtype>
void Blob<Dtype>::set_delta(){
float scale_factor = TERNARY_DELTA * 1.0 / 10; //delta = 0.7
Dtype delta = (Dtype) scale_factor * this->asum_data() / this->count(); // 0.7*(E|W_i|)/num
delta = (delta <= 100) ? delta : 100;
delta = (delta >= -100) ? delta : -100;
this->delta_ = delta;
}
template <typename Dtype>
void Blob<Dtype>::set_delta(Dtype delta){
delta = (delta <= 100) ? delta : 100;
delta = (delta >= -100) ? delta : -100;
this->delta_ = delta;
}
Implement
我借用一张图
步骤3~5,其中第5步代码在上面:
template <typename Dtype>
void Blob<Dtype>::ternarize_data(Phase phase){
if(phase == RUN){
// if(DEBUG) print_head();
//LOG(INFO) << "RUN phase...";
// caffe_sleep(3);
return; // do nothing for the running phase
}else if(phase == TRAIN){
//LOG(INFO) << "TRAIN phase ...";
// caffe_sleep(3);
}else{
//LOG(INFO) << "TEST phase ...";
// caffe_sleep(3);
}
// const Dtype delta = 0; // default value;
// const Dtype delta = (Dtype) 0.8 * this->asum_data() / this->count();
this->set_delta(); //defualt 0.7*(E|W_i|)/num or set by user
const Dtype delta = this->get_delta();
Dtype alpha = 1;
if (!data_) { return; }
switch (data_->head()) {
case SyncedMemory::HEAD_AT_CPU:
{
caffe_cpu_ternary<Dtype>(this->count(), delta, this->cpu_data(), this->mutable_cpu_binary()); //quantized weight to ternary
alpha = caffe_cpu_dot(this->count(), this->cpu_binary(), this->cpu_data()); //scale-alpha: (E |W_i|) i belong to I_delta
alpha /= caffe_cpu_dot(this->count(), this->cpu_binary(), this->cpu_binary()); //(1/num_binary)*alpha
caffe_cpu_scale(this->count(), alpha, this->cpu_binary(), this->mutable_cpu_binary());
// this->set_alpha(alpha);
}
return;
case SyncedMemory::HEAD_AT_GPU:
case SyncedMemory::SYNCED:
#ifndef CPU_ONLY
{
caffe_gpu_ternary<Dtype>(this->count(), delta, this->gpu_data(), this->mutable_gpu_binary());
Dtype* pa = new Dtype(0);
caffe_gpu_dot(this->count(), this->gpu_binary(), this->gpu_data(), pa);
Dtype* pb = new Dtype(0);
caffe_gpu_dot(this->count(), this->gpu_binary(), this->gpu_binary(), pb);
alpha = (*pa) / ((*pb) + 1e-6);
this->set_alpha(alpha);
caffe_gpu_scale(this->count(), alpha, this->gpu_binary(), this->mutable_gpu_binary());
// this->set_alpha((Dtype)1);
// LOG(INFO) << "alpha = " << alpha;
// caffe_sleep(3);
}
return;
#else
NO_GPU;
#endif
case SyncedMemory::UNINITIALIZED:
return;
default:
LOG(FATAL) << "Unknown SyncedMemory head state: " << data_->head();
}
}
步骤6~7,其中在第6步作者在caffe-twns直接采用传统caffe的方法,而$Z=XW\approx X(\alpha W^t)=(\alpha X)\bigoplus W^t $更偏向与在硬件加速的优化(因为本身在caffe-twns的ternary就采用float或者double,并且用blas或cudnn加速也无法直接跳过0值):
//conv_layer.cpp
template <typename Dtype>
void ConvolutionLayer<Dtype>::Forward_cpu(const vector<Blob<Dtype>*>& bottom,
const vector<Blob<Dtype>*>& top) {
// const Dtype* weight = this->blobs_[0]->cpu_data();
if(BINARY){
this->blobs_[0]->binarize_data();
}
if(TERNARY){
this->blobs_[0]->ternarize_data(this->phase_); //quantized from blob[0] to ternary sand stored in cpu_binary()
/*
Dtype alpha = (Dtype) this->blobs_[0]->get_alpha();
for(int i=0; i<bottom.size(); i++){
Blob<Dtype>* blob = bottom[i];
caffe_cpu_scale(blob->count(), alpha, blob->cpu_data(), blob->mutable_cpu_data());
}
*/
}
const Dtype* weight = (BINARY || TERNARY) ? this->blobs_[0]->cpu_binary() : this->blobs_[0]->cpu_data();
...
}
步骤11~19,weight的Update是在full precision上,而计算gradient则是用ternary weight:
//conv_layer.cpp
template <typename Dtype>
void ConvolutionLayer<Dtype>::Backward_cpu(const vector<Blob<Dtype>*>& top,
const vector<bool>& propagate_down, const vector<Blob<Dtype>*>& bottom) {
const Dtype* weight = this->blobs_[0]->cpu_data();
Dtype* weight_diff = this->blobs_[0]->mutable_cpu_diff();
for (int i = 0; i < top.size(); ++i) {
...
if (this->param_propagate_down_[0] || propagate_down[i]) {
for (int n = 0; n < this->num_; ++n) {
// gradient w.r.t. weight. Note that we will accumulate diffs.
if (this->param_propagate_down_[0]) {
this->weight_cpu_gemm(bottom_data + n * this->bottom_dim_,
top_diff + n * this->top_dim_, weight_diff);
}
// gradient w.r.t. bottom data, if necessary.
if (propagate_down[i]) {
this->backward_cpu_gemm(top_diff + n * this->top_dim_, weight,
bottom_diff + n * this->bottom_dim_);
}
}
}
}
}