mirror of
https://github.com/saymrwulf/pytorch.git
synced 2026-05-15 21:00:47 +00:00
106b0b6a62
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/40760 Add op to create a quant scheme. Test Plan: buck test mode/opt caffe2/caffe2/quantization/server:int8_quant_scheme_blob_fill_test {F241838981} Reviewed By: csummersea Differential Revision: D22228154 fbshipit-source-id: 1b7a02c06937c68e2fcccf77eb10a965300ed732
33 lines
1.1 KiB
C++
33 lines
1.1 KiB
C++
// Copyright 2004-present Facebook. All Rights Reserved.
|
|
|
|
#pragma once
|
|
#include "caffe2/quantization/server/caffe2_dnnlowp_utils.h"
|
|
#include "caffe2/quantization/server/dnnlowp.h"
|
|
#include "caffe2/quantization/server/int8_gen_quant_params.h"
|
|
|
|
namespace caffe2 {
|
|
using namespace std;
|
|
|
|
template <class Context, class Engine = DefaultEngine>
|
|
class Int8QuantSchemeBlobFillOp final : public Operator<Context> {
|
|
public:
|
|
USE_OPERATOR_CONTEXT_FUNCTIONS;
|
|
Int8QuantSchemeBlobFillOp(const OperatorDef& operator_def, Workspace* ws)
|
|
: Operator<Context>(operator_def, ws) {}
|
|
bool RunOnDevice() override {
|
|
std::string quantization_kind =
|
|
this->template GetSingleArgument<std::string>(
|
|
"quantization_kind", "MIN_MAX_QUANTIZATION");
|
|
bool preserve_sparsity =
|
|
this->template GetSingleArgument<bool>("preserve_sparsity", false);
|
|
|
|
auto* output_qscheme =
|
|
this->template Output<unique_ptr<Int8QuantSchemeBlob>>(0);
|
|
output_qscheme->reset(
|
|
new Int8QuantSchemeBlob(quantization_kind, preserve_sparsity));
|
|
return true;
|
|
}
|
|
|
|
}; // class Int8QuantSchemeBlobFillOp
|
|
|
|
} // namespace caffe2
|