blob: 92008b43e07b4506c0e39ff28e72cb631fe4309f [file] [log] [blame]
#include "caffe2/opt/backend_transformer_base.h"
#include "caffe2/onnx/onnx_exporter.h"
#include "caffe2/utils/proto_utils.h"
namespace caffe2 {
// Populate 'net_pos' argument for any ops that don't already have it. 'net_pos'
// we populate here starts after the max 'net_pos' value we encountered.
void BackendTransformerBase::annotateOpIndex(NetDef* net) {
// find the max net_pos that we have so far.
int i = -1;
for (const auto& op : net->op()) {
ArgumentHelper helper(op);
int old_index = helper.GetSingleArgument(op, kNetPos, -1);
i = std::max(i, old_index);
}
// populate net_pos for any op that doesn't already have it.
for (auto& op : *(net->mutable_op())) {
if (!ArgumentHelper::HasArgument(op, kNetPos)) {
AddArgument(kNetPos, ++i, &op);
}
}
}
std::string BackendTransformerBase::getModelId(const NetDef& net) {
static std::atomic<size_t> seq_id{0};
std::string model_id;
for (const auto& arg : net.arg()) {
if (arg.name() == kModelId) {
if (arg.has_s()) {
model_id = arg.s();
} else if (arg.has_i()) {
model_id = c10::to_string(arg.i());
}
break;
}
}
if (model_id.empty()) {
model_id = "unnamed_" + c10::to_string(seq_id++);
}
return model_id;
}
TensorProto wrapShapeInfoIntoTensorProto(
const std::string& name,
const ShapeInfo& shape_info) {
TensorProto t;
t.set_name(name);
t.set_data_type(shape_info.shape.data_type());
for (const auto i : shape_info.shape.dims()) {
t.add_dims(i);
}
for (const auto& dimType : shape_info.getDimType()) {
t.add_int32_data(static_cast<int32_t>(dimType));
}
return t;
}
QTensorProto wrapShapeInfoIntoQTensorProto(
const std::string& name,
const ShapeInfo& shape_info) {
QTensorProto t;
CAFFE_ENFORCE(
shape_info.is_quantized == true,
"Only quantized shapeinfo can be extracted into QTensor!");
t.set_name(name);
t.set_data_type(shape_info.shape.data_type());
t.set_axis(shape_info.q_info.axis);
t.set_is_multiparam(true);
for (const auto i : shape_info.q_info.scale) {
t.add_scales(i);
}
t.set_scale(1.0);
for (const auto i : shape_info.q_info.offset) {
t.add_biases(i);
}
t.set_bias(0.0);
// precision and is_signed is not used in onnxifi workflow, but it is required
// field
t.set_precision(0);
// NOLINTNEXTLINE(modernize-use-bool-literals)
t.set_is_signed(0);
for (const auto i : shape_info.shape.dims()) {
t.add_dims(i);
}
for (const auto& dimType : shape_info.getDimType()) {
t.add_data(static_cast<int32_t>(dimType));
}
return t;
}
ShapeInfoMap BackendTransformerBase::ssaRewriteAndMapNames(
Workspace* ws,
NetDef* pred_net,
const ShapeInfoMap& input_shape_hints) {
input_mapping_ = onnx::SsaRewrite(nullptr, pred_net);
// Annote the ops with net position
annotateOpIndex(pred_net);
// Since we are going to create a mapped workspace, we need to make sure that
// the parent workspace has the mapped blob names. If the blobs don't exist
// (usually such blobs are input tensor names), we exclude them from mapping.
std::vector<std::string> exclude_mapping;
for (const auto& kv : input_mapping_) {
if (!ws->HasBlob(kv.second)) {
exclude_mapping.emplace_back(kv.first);
}
}
for (const auto& i : exclude_mapping) {
input_mapping_.erase(i);
}
ShapeInfoMap shape_hints_mapped;
for (const auto& kv : input_shape_hints) {
shape_hints_mapped.emplace(kv.first, kv.second);
}
return shape_hints_mapped;
}
ShapeInfoMap BackendTransformerBase::inferShapes(
Workspace* ws,
NetDef* pred_net,
const ShapeInfoMap& shape_hints_mapped,
const BoundShapeSpec& spec) {
ShapeInfoMap shape_map;
// Populate shapes from workplace
const std::vector<std::string> ws_blobs = ws->Blobs();
for (const auto& s : ws_blobs) {
auto shape_info = getShapeInfoFromBlob(ws->GetBlob(s));
if (shape_info.dimTypeIsSet()) {
shape_map.emplace(s, shape_info);
}
}
for (const auto& s : shape_hints_mapped) {
shape_map.insert(s);
}
auto eng = BoundShapeInferencerRegistry()->Create("C10", spec);
eng->InferBoundShapeAndType(*pred_net, shape_map, ws);
const auto& out_map = eng->shape_info();
shape_map.clear();
for (const auto& kv : out_map) {
shape_map.emplace(
std::piecewise_construct,
std::forward_as_tuple(kv.first),
std::forward_as_tuple(
kv.second.getDimType(),
kv.second.shape,
kv.second.is_quantized,
kv.second.q_info));
}
return shape_map;
}
void BackendTransformerBase::addShapeToNet(
NetDef& shape_net,
const ShapeInfoMap& shape_hints) const {
auto* shape_arg = shape_net.add_arg();
auto* qshape_arg = shape_net.add_arg();
shape_arg->set_name("shape_info");
qshape_arg->set_name("qshape_info");
for (const auto& kv : shape_hints) {
if (!kv.second.is_quantized) {
auto t = wrapShapeInfoIntoTensorProto(kv.first, kv.second);
shape_arg->mutable_tensors()->Add()->CopyFrom(t);
} else {
auto t = wrapShapeInfoIntoQTensorProto(kv.first, kv.second);
qshape_arg->mutable_qtensors()->Add()->CopyFrom(t);
}
}
}
void BackendTransformerBase::dumpNet(
const NetDef& pred_net,
const ShapeInfoMap& shape_hints,
const std::string& fname) const {
NetDef shape_net(pred_net);
addShapeToNet(shape_net, shape_hints);
WriteProtoToTextFile(shape_net, fname, false);
}
} // namespace caffe2