最近有了突如其来的想法,如何把caffe的变得更小更快。后来翻到Nvidia开发caffe-float16,同时也看到它的论文。看完大致了解一番后,就做一下记录。

该工作的目标是,减少网络的所需的内存大小和提升网络的 inference(推理)速度。nvidia通过才用自己开发的 float16 半精度 cuda_fp16.h 数据类型,在forwardbackward propagation中代替 float 32 bits的单精度数据类型。因此,在降低网络的数据的 precision 时候,导致产生了网络 accuracy 降低和 gradient 消失无法收敛的问题。当然,我在这里并不想重复的写出文中所有的点(因为其中总体的idea在量化quantization 方面是“general” 的),仅对该工作我觉得特有的点或感兴趣的点进行简述。

Mixed Precision

Mixed Precision Training —— caffe-float16
caffe-float16 中的Blob重写,改为datadiff分别用不同的数据类型表示,这可以选着你所需的精确的数据类型:

//blob.hpp
 protected:
  Blob(Type data_type, Type diff_type)
      : data_tensor_(make_shared<Tensor>(data_type)),
        diff_tensor_(make_shared<Tensor>(diff_type)),
        count_(0) {}

Master-Weights(F32)-->float2half的实现就是每次this->blobs_[0]->template gpu_data<Ftype>(); 中做一次类型转换:

//conv_layer.cu
const Ftype* weight = this->blobs_[0]->template gpu_data<Ftype>();
//blob.hpp
  template<typename Dtype>
  const Dtype* gpu_data() const {
    convert_data(tp<Dtype>());
    return static_cast<const Dtype*>(data_tensor_->synced_mem()->gpu_data());
  }
  
  void convert_data(Type new_data_type) const {
    data_tensor_->convert(new_data_type);
  }
//tensor.cpp
void Tensor::convert(Type new_type) {
  if (new_type == type_) {
    return;
  }
  const shared_ptr<SyncedMemory>& current_mem = synced_mem();
  shared_ptr<SyncedMemory>& new_mem = synced_arrays_->at(new_type);

  if (!new_mem || !new_mem->is_valid()) {
    const std::size_t new_cap = even(count_) * tsize(new_type);
    if (!new_mem || new_mem->size() != new_cap) {
      new_mem = make_shared<SyncedMemory>(new_cap);
    }
    const bool data_gpu = Caffe::mode() == Caffe::GPU;
    if (current_mem->head() != SyncedMemory::UNINITIALIZED) {
      copy_helper(data_gpu, count_,
          data_gpu ? current_mem->gpu_data() : current_mem->cpu_data(),
          type_,
          data_gpu ? new_mem->mutable_gpu_data() : new_mem->mutable_cpu_data(),
          new_type);
    }
  } // we just trust its current status otherwise
  type_ = new_type;
  new_mem->validate();
}

神经网络的 quantization 一般可分 activation、weight 部分,当然也存在继续对不同类型 layerweghit 进行 quantization 的。而 Nvidia 则提出了 gradient 也是要 quantization 。上图是文中的整个方法的流程图,为了防止用无法拟合,采用全精度的 flaot32 来保存完整的权重信息(其他文章又叫 full precision shadow weight ),每次 forward 是都做copyround/quantization 。 这是有两个原因:

  1. 因为 gradient x learning rate < (2^{-24}) ,小于float16 范围,导致梯度消失无法更新。Mixed Precision Training —— caffe-float16
    2.由于浮点型的特性,相加时会进行小数点对齐(即对其 exponent)。由于float16 表示的weightfloat16表示的 gradient 相差2048倍(因为float16mantissa 只有 10bits,有右移超过11bits ,即2048倍),则 gradient 变成0。float16 各个部分:
    Mixed Precision Training —— caffe-float16
    除非指数位全是0,否则就会假定隐藏的起始位是1。因此只有10位 mantissa在内存中被显示出来,而总精度是11位。据IEEE 754的说法,虽然尾数只有10位,但是尾数精度是11位的(log10(211) ≈ 3.311 十进制数).

Weight Update,会对diff进行类型转换

//blob.hpp

// The "update" method is used for parameter blobs in a Net, which are stored
// as TBlob<float> or TBlob<double> -- hence we do not define it for
// TBlob<int> or TBlob<unsigned int>.
void Blob::Update() {
  convert_diff(data_type());  // align data&diff types
  shared_ptr<SyncedMemory>& data_mem = data_tensor_->mutable_synced_mem();
  const shared_ptr<SyncedMemory>& diff_mem = diff_tensor_->synced_mem();
  // We will perform update based on where the data is located.
  switch (data_mem->head()) {
  case SyncedMemory::HEAD_AT_CPU:
    // perform computation on CPU
    cpu_axpy(count_, data_type(), -1.F,
        diff_mem->cpu_data(), data_mem->mutable_cpu_data());
    break;
  case SyncedMemory::HEAD_AT_GPU:
  case SyncedMemory::SYNCED:
#ifndef CPU_ONLY
    gpu_axpy(count_, data_type(), -1.F,
        diff_mem->gpu_data(), data_mem->mutable_gpu_data());
#else
    NO_GPU;
#endif
    break;
    default:
    LOG(FATAL) << "Syncedmem not initialized.";
  }
  CHECK(is_current_data_valid());
  CHECK(is_current_diff_valid());
}

Lose Scaling

从上面float16 各个部分位宽可以得到,float16 可以表示的范围是([2^{-24},2^{15}])(exponent表示范围是([2^{-14},2^{15}]),其中 mantissa10bits)。但是 activationgradient的分布却在 ([2^{-60},2^{-10}]),在float16 中有非常大的表示范围并没用,同时导致大多数的activation gradient变成0。因此,对activation gradientforward后,backward propagation前做scaling/shift 。并且,在链式法则backward propagation 中的所有activation gradient按想用的量进行scaling
Mixed Precision Training —— caffe-float16
具体操作(因为activation gradientscaling,那么也要对learning rateweight_decayscaling):

#caffe train_val.prototxt
#To sfift gradients dE/dX we will scale up the loss function by constant (e.g. by 1000):
layer {
	type: "SoftMaxWithLoss"
	loo_weight: 1000.
}
#and adjust learning rate and weights decay accordingly

base_lr: 0.00001 #(original value is 0.01, 0.01 / 1000)
weight_decay: 0.5  #(original value is 0.0005, 0.5 * 1000)

其中decay_weight公式为:

[omega_i leftarrow omega_i - eta{{partial E}over partial omega_i} - etalambda omega_i
]

而在softmax_loss_layer.cu的实现为:

template <typename Ftype, typename Btype>
void SoftmaxWithLossLayer<Ftype, Btype>::Backward_gpu(const vector<Blob*>& top,
    const vector<bool>& propagate_down, const vector<Blob*>& bottom) {
	...
    float loss_weight = float(top[0]->cpu_diff<Btype>()[0]) /
                              get_normalizer(normalization_, valid_count);
    if (this->parent_net() != NULL) {
      loss_weight *= this->parent_net()->global_grad_scale();
    }
    caffe_gpu_scal<Btype>(prob_->count(), loss_weight , bottom_diff);
  }
}

FP16 Master Weight Storage

在该论文之外,Nvidia还考虑避免每次foward都复制权重,用float16进行权重更新的问题。

  • 最核心一点就是避免gradient (eta{{partial E}over partial omega_i}=etaDelta omega_i)消失。

那么Nvidia提出对momentum SGD 进行改进

  1. Compute momentum (H) : (H(t+1)=m*H(t)-lambda Delta W(t))
  2. Update wights with (H): (W(t+1)=W(t)+H(t+1))

假设(lambda)为常数,把式①展开:

[H(t+1)=m*H(t)-lambda Delta W(t)=m*(m*H(t-1)-lambda Delta W(t-1))-lambda Delta W(t)
]

[=-lambda [Delta W(t)+mDelta W(t-1)+m^2Delta W(t-2)+m^kDelta W(t-k)+...]
]

因此新的公式:

  1. Compute momentum (H) : (H(t+1)=m*H(t)-color{#F00}{cancel{lambda}}Delta W(t))
  2. Update wights with (H): (W(t+1)=W(t)+color{#F00}{lambda} H(t+1))

这样可以避免(H)(lambda Delta W(t))消失时,momentum不断的消失。因为新的公式避免了(Delta W(t))的消失,而且momentum会不断更新。
ps:这里Nvidia解释是 Moment works as average of gradients.

Nvidia的总结

懒癌犯了- -!!
Mixed Precision Training —— caffe-float16