#ifndef QUANT_DECODE_OP_H_
|
#define QUANT_DECODE_OP_H_
|
|
#include "caffe2/core/context.h"
|
#include "caffe2/core/operator.h"
|
#include "caffe2/core/tensor.h"
|
#include <c10/util/typeid.h>
|
|
namespace caffe2 {
|
|
namespace {
|
|
template <class CodebookT, class CodeT>
|
void Decode(
|
const Tensor& codebook,
|
const Tensor& codes,
|
/* optional */ const Tensor* const decoded_grad,
|
Tensor* const output,
|
bool resizeOnly) {
|
CAFFE_ENFORCE(codebook.IsType<CodebookT>());
|
|
auto* cb_ptr = codebook.data<CodebookT>();
|
int cb_size = codebook.numel();
|
|
CAFFE_ENFORCE(codes.IsType<CodeT>());
|
auto* code_ptr = codes.data<CodeT>();
|
|
if (decoded_grad == nullptr) {
|
// Forward pass: decode and store codebook values in output.
|
output->ResizeLike(codes);
|
auto* out_ptr = output->template mutable_data<CodebookT>();
|
if (resizeOnly) {
|
return;
|
}
|
|
int sz = output->numel();
|
for (int i = 0; i < sz; i++) {
|
DCHECK_LE(*code_ptr, cb_size);
|
*out_ptr++ = cb_ptr[*code_ptr++];
|
}
|
} else {
|
// Backward pass: decode and accumulate gradient w.r.t. codebook values.
|
CAFFE_ENFORCE_EQ(codes.numel(), decoded_grad->numel());
|
auto* gradient_ptr = decoded_grad->data<CodebookT>();
|
auto* const gradient_end = gradient_ptr + decoded_grad->numel();
|
|
CAFFE_ENFORCE_EQ(cb_size, output->numel());
|
auto* out_ptr = output->template mutable_data<CodebookT>();
|
while (gradient_ptr < gradient_end) {
|
DCHECK_LE(*code_ptr, cb_size);
|
out_ptr[*code_ptr++] += *gradient_ptr++;
|
}
|
}
|
}
|
|
#define REGISTER_DECODER(codebookType, codesType) \
|
{ \
|
{TypeMeta::Id<codebookType>(), TypeMeta::Id<codesType>()}, \
|
[](const Tensor& codebook_, \
|
const Tensor& codes_, \
|
const Tensor* gradient_, \
|
Tensor* outDecoded_, \
|
bool resizeOnly_) { \
|
Decode<codebookType, codesType>( \
|
codebook_, codes_, gradient_, outDecoded_, resizeOnly_); \
|
} \
|
}
|
|
inline void DecodeGeneral(
|
const Tensor& codebook,
|
const Tensor& codes,
|
const Tensor* gradient,
|
Tensor* outDecoded,
|
bool resizeOnly) {
|
const static std::map<
|
std::pair<TypeIdentifier, TypeIdentifier>,
|
std::function<void(
|
const Tensor& codebook,
|
const Tensor& codes,
|
const Tensor* gradient,
|
Tensor* outDecoded,
|
bool resizeOnly)>>
|
gDecoderMapper = {REGISTER_DECODER(float, uint8_t),
|
REGISTER_DECODER(float, uint16_t),
|
REGISTER_DECODER(float, int32_t)};
|
|
gDecoderMapper.at({codebook.dtype().id(), codes.dtype().id()})(
|
codebook, codes, gradient, outDecoded, resizeOnly);
|
}
|
|
} // namespace
|
|
// Decode tensors based on given codebook,
|
// The codebook is generated by model_quantize.py
|
|
enum class QuantDecodeRunTy {
|
RUN_ALWAYS,
|
RUN_ONCE,
|
};
|
|
template <QuantDecodeRunTy QuantDecodeRun>
|
class QuantDecodeOp final : public Operator<CPUContext> {
|
public:
|
USE_OPERATOR_FUNCTIONS(CPUContext);
|
template <class... Args>
|
explicit QuantDecodeOp(Args&&... args)
|
: Operator<CPUContext>(std::forward<Args>(args)...) {}
|
|
~QuantDecodeOp() {}
|
|
bool RunOnDevice() override {
|
CAFFE_ENFORCE_GT(InputSize(), 1);
|
// first input is the codebook
|
CAFFE_ENFORCE_EQ(InputSize(), OutputSize() + 1);
|
|
const auto& codebook = Input(0);
|
CAFFE_ENFORCE(codebook.template IsType<float>(), codebook.dtype().name());
|
|
for (int i = 0; i < OutputSize(); i++) {
|
auto& ci = Input(i + 1);
|
auto* co = Output(i);
|
|
DecodeGeneral(
|
codebook,
|
ci,
|
nullptr,
|
co,
|
/*resizeOnly=*/QuantDecodeRun == QuantDecodeRunTy::RUN_ONCE &&
|
hasRun_);
|
}
|
hasRun_ = true;
|
return true;
|
}
|
|
private:
|
bool hasRun_{false};
|
};
|
|
class QuantDecodeGradientOp final : public Operator<CPUContext> {
|
public:
|
USE_OPERATOR_FUNCTIONS(CPUContext);
|
template <class... Args>
|
explicit QuantDecodeGradientOp(Args&&... args)
|
: Operator<CPUContext>(std::forward<Args>(args)...) {}
|
~QuantDecodeGradientOp() {}
|
|
bool RunOnDevice() override {
|
// Inputs: 1 codebook, n tensors of codes, and n corresponding gradients.
|
CAFFE_ENFORCE(InputSize() >= 3 && InputSize() % 2 == 1);
|
const int num_code_tensors = (InputSize() - 1) / 2;
|
CAFFE_ENFORCE_EQ(OutputSize(), 1);
|
|
const auto& codebook = Input(0);
|
CAFFE_ENFORCE(codebook.template IsType<float>(), codebook.dtype().name());
|
|
auto* gradient = Output(0, codebook.sizes(), at::dtype<float>());
|
auto* gradient_ptr = gradient->template mutable_data<float>();
|
std::fill(gradient_ptr, gradient_ptr + gradient->numel(), 0);
|
|
for (int i = 0; i < num_code_tensors; i++) {
|
auto& codes_i = Input(i + 1);
|
auto& output_gradient_i = Input(i + num_code_tensors + 1);
|
DecodeGeneral(codebook, codes_i, &output_gradient_i, gradient, false);
|
}
|
return true;
|
}
|
};
|
|
} // namespace caffe2
|
#endif // QUANT_DECODE_OP_H_
|
