Mixed Precision Training —— caffe-float16
简介
最近有了突如其来的想法,如何把caffe的变得更小更快。后来翻到Nvidia开发caffe-float16,同时也看到它的论文。看完大致了解一番后,就做一下记录。
该工作的目标是,减少网络的所需的内存大小和提升网络的 inference(推理)速度。nvidia通过才用自己开发的 float16 半精度 cuda_fp16.h 数据类型,在forward和backward propagation中代替 float 32 bits的单精度数据类型。因此,在降低网络的数据的 precision 时候,导致产生了网络 accuracy 降低和 gradient 消失无法收敛的问题。当然,我在这里并不想重复的写出文中所有的点(因为其中总体的idea在量化quantization 方面是“general” 的),仅对该工作我觉得特有的点或感兴趣的点进行简述。
Mixed Precision
在caffe-float16 中的Blob重写,改为data和diff分别用不同的数据类型表示,这可以选着你所需的精确的数据类型:
//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 部分,当然也存在继续对不同类型 layer 的 weghit 进行 quantization 的。而 Nvidia 则提出了 gradient 也是要 quantization 。上图是文中的整个方法的流程图,为了防止用无法拟合,采用全精度的 flaot32 来保存完整的权重信息(其他文章又叫 full precision shadow weight ),每次 forward 是都做copy 和 round/quantization 。 这是有两个原因:
- 因为 gradient x learning rate < \(2^{-24}\) ,小于float16 范围,导致梯度消失无法更新。
2.由于浮点型的特性,相加时会进行小数点对齐(即对其 exponent)。由于float16 表示的weight 与 float16表示的 gradient 相差2048倍(因为float16 的 mantissa 只有 10bits,有右移超过11bits ,即2048倍),则 gradient 变成0。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}]\),其中 mantissa是10bits)。但是 activation的 gradient的分布却在 \([2^{-60},2^{-10}]\),在float16 中有非常大的表示范围并没用,同时导致大多数的activation gradient变成0。因此,对activation gradient在forward后,backward propagation前做scaling/shift 。并且,在链式法则backward propagation 中的所有activation gradient按想用的量进行scaling。
具体操作(因为activation gradient做scaling,那么也要对learning rate和weight_decay做scaling):
#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公式为:
而在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}=\eta\Delta \omega_i\)消失。
那么Nvidia提出对momentum SGD 进行改进
- Compute momentum \(H\) : \(H(t+1)=m*H(t)-\lambda \Delta W(t)\)
- Update wights with \(H\): \(W(t+1)=W(t)+H(t+1)\)
假设\(\lambda\)为常数,把式①展开:
因此新的公式:
- Compute momentum \(H\) : \(H(t+1)=m*H(t)-\color{#F00}{\cancel{\lambda}}\Delta W(t)\)
- 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的总结
懒癌犯了- -!!