mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
4b0fc5200b
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/13689 Now that typeid.h lives in c10/util, the include paths should reflect that. Reviewed By: ezyang Differential Revision: D12912237 fbshipit-source-id: e54225f049f690de77cb6d5f417994b211a6e1fb
68 lines
2.1 KiB
C++
68 lines
2.1 KiB
C++
#include "quant_decode_op.h"
|
|
#include <stdint.h>
|
|
#include "caffe2/core/tensor.h"
|
|
#include <c10/util/typeid.h>
|
|
|
|
namespace caffe2 {
|
|
|
|
REGISTER_CPU_OPERATOR(QuantDecode, QuantDecodeOp<QuantDecodeRunTy::RUN_ALWAYS>);
|
|
REGISTER_CPU_GRADIENT_OPERATOR(QuantDecodeGradient, QuantDecodeGradientOp);
|
|
#ifdef CAFFE2_USE_MPSCNN
|
|
REGISTER_CPU_OPERATOR(
|
|
MPSCNNQuantDecode,
|
|
QuantDecodeOp<QuantDecodeRunTy::RUN_ONCE>);
|
|
#endif
|
|
|
|
OPERATOR_SCHEMA(QuantDecode)
|
|
.NumInputsOutputs([](int in, int out) { return in > 1 && out + 1 == in; })
|
|
.SetDoc(R"DOC(
|
|
Decode inputs using codebook. This is a general LUT operator that returns
|
|
tensors with values from codebook (input 0) based on given indices in
|
|
codes (input 1 ~ n).
|
|
|
|
|
|
Example:
|
|
|
|
|
|
Input:
|
|
codebook = [1.5, 2.5, 3.5]
|
|
codes_0 = [0, 1, 1, 2]
|
|
codes_1 = [2, 0, 0]
|
|
|
|
|
|
Output:
|
|
decoded_0 = [1.5, 2.5, 2.5, 3.5]
|
|
decoded_1 = [3.5, 1.5, 1.5]
|
|
)DOC")
|
|
.Input(0, "codebook", "Codebook in 1d tensor (float)")
|
|
.Input(1, "codes_0", "Encoded codes 0 (uint8/uint16/int32)")
|
|
.Input(2, "codes_1", "Encoded codes 1 if existed (uint8/uint16/int32)")
|
|
.Input(3, "codes_n", "Encoded codes n if existed (uint8/uint16/int32)")
|
|
.Output(0, "decoded_0", "Decoded tensor for codes_0 (float)")
|
|
.Output(1, "decoded_1", "Decoded tensor for codes_1 (float)")
|
|
.Output(2, "decoded_n", "Decoded tensor for codes_n (float)");
|
|
|
|
GRADIENT_OPERATOR_SCHEMA(QuantDecodeGradient)
|
|
.NumInputs([](int in) { return in >= 3 && in % 2 == 1; })
|
|
.NumOutputs(1);
|
|
|
|
class GetQuantDecodeGradient : public GradientMakerBase {
|
|
using GradientMakerBase::GradientMakerBase;
|
|
vector<OperatorDef> GetGradientDefs() override {
|
|
CAFFE_ENFORCE_EQ(Def().input_size(), Def().output_size() + 1);
|
|
vector<string> gradient_op_inputs;
|
|
for (int i = 0; i < Def().input_size(); i++) {
|
|
gradient_op_inputs.push_back(I(i));
|
|
}
|
|
for (int i = 0; i < Def().output_size(); i++) {
|
|
gradient_op_inputs.push_back(GO(i));
|
|
}
|
|
return SingleGradientDef(
|
|
"QuantDecodeGradient", "", gradient_op_inputs, vector<string>{GI(0)});
|
|
}
|
|
};
|
|
|
|
REGISTER_GRADIENT(QuantDecode, GetQuantDecodeGradient);
|
|
|
|
} // namespace caffe2
|