caffe2/predictor/transforms.cc - platform/external/pytorch - Git at Google

 #include "caffe2/predictor/transforms.h"
 #include "caffe2/onnx/onnx_exporter.h"
 #include "caffe2/utils/proto_utils.h"

 #include <unordered_set>

 namespace caffe2 {

 namespace {
 bool HasInput(const string& blob, const OperatorDef& op) {
   for (const auto& in : op.input()) {
     if (blob == in) {
       return true;
     }
   }
   return false;
 }

 bool HasOutput(const string& blob, const OperatorDef& op) {
   for (const auto& out : op.output()) {
     if (blob == out) {
       return true;
     }
   }
   return false;
 }

 void RewriteSubnetsForIfOp(
     const string& from,
     const string& to,
     OperatorDef* op) {
   ArgumentHelper helper(*op);
   Argument *then_arg = nullptr, *else_arg = nullptr;

   std::map<std::string, std::string> oldname_to_newname;
   oldname_to_newname[from] = to;

   if (helper.HasSingleArgumentOfType<NetDef>("then_net")) {
     then_arg = GetMutableArgument("then_net", false, op);
     onnx::rewriteSubnet(then_arg, oldname_to_newname);
   }
   if (helper.HasSingleArgumentOfType<NetDef>("else_net")) {
     else_arg = GetMutableArgument("else_net", false, op);
     onnx::rewriteSubnet(else_arg, oldname_to_newname);
   }
 }

 void RenameInputs(
     const string& from,
     const string& to,
     OperatorDef* def,
     int op_idx,
     std::unordered_map<std::string, std::unordered_set<int>>& children) {
   VLOG(2) << "RenameInputs (from=" << from << ", to=" << to << ", "
           << def->DebugString() << ")";
   for (int i = 0; i < def->input_size(); i++) {
     if (def->input(i) == from) {
       *def->mutable_input(i) = to;
       children[from].erase(op_idx);
       children[to].insert(op_idx);
     }
   }
   // Rename inputs in the subnets of If/AsyncIf op
   if (def->type() == "If" || def->type() == "AsyncIf") {
     RewriteSubnetsForIfOp(from, to, def);
   }
 }

 void RenameOutputs(
     const string& from,
     const string& to,
     OperatorDef* def,
     int op_idx,
     std::unordered_map<std::string, std::unordered_set<int>>& parents) {
   VLOG(2) << "RenameOutputs (from=" << from << ", to=" << to << ", "
           << def->DebugString() << ")";
   for (string& output : *def->mutable_output()) {
     if (output == from) {
       output = to;
       parents[from].erase(op_idx);
       parents[to].insert(op_idx);
     }
   }
   // Rename outputs in the subnets of If/AsyncIf op
   if (def->type() == "If" || def->type() == "AsyncIf") {
     RewriteSubnetsForIfOp(from, to, def);
   }
 }

 void RenameInputsInChildren(
     const string& from,
     const string& to,
     caffe2::NetDef* net,
     std::unordered_map<std::string, std::unordered_set<int>>& children) {
   VLOG(2) << "RenameInputsInChildren (from=" << from << ", to=" << to << ")";
   if (children.count(from) == 0) {
     return;
   }

   // make an temporary copy here because we're going to modify children
   for (int child : std::unordered_set<int>(children[from])) {
     RenameInputs(from, to, net->mutable_op(child), child, children);
   }
 }

 void RenameOutputInParents(
     const std::string& from,
     const std::string& to,
     caffe2::NetDef* net,
     std::unordered_map<std::string, std::unordered_set<int>>& parents) {
   VLOG(2) << "RenameOutputInParents (from=" << from << ", to=" << to << ")";
   if (parents.count(from) == 0) {
     return;
   }

   // make an temporary copy here because we're going to modify parents
   for (int parent : std::unordered_set<int>(parents[from])) {
     RenameOutputs(from, to, net->mutable_op(parent), parent, parents);
   }
 }

 bool FoundOpCandidate(
     const OperatorDef* op,
     int op_idx,
     const std::string& op_type,
     const std::unordered_set<std::string>& inputs,
     const std::unordered_set<std::string>& outputs,
     const std::unordered_map<std::string, std::unordered_set<int>>& parents,
     const std::unordered_map<std::string, std::unordered_set<int>>& children) {
   if (op->type() != op_type) {
     VLOG(2) << "InplaceOps(" << op_type << ") skipping op: \n"
             << op->DebugString();
     return false;
   }
   if (op->input_size() != 1 || op->output_size() != 1) {
     VLOG(2) << "InplaceOps(" << op_type
             << ") only supports ops with exactly 1 output "
             << "and exactly 1 input. Skipping op: \n"
             << op->DebugString();
     return false;
   }

   // use actual copy because op->input/output may change
   const std::string in = op->input(0);
   const std::string out = op->output(0);

   if (in == out) {
     // This case can still exist when in/out is in the predict_net's outputs.
     // The op is an inplace op already.
     return false;
   }

   // The following is to handle the special cases of inputs being overwritten
   // by ops in the net and then appear in outputs of the net
   if (outputs.count(out) == 0) {
     // Propagate input downwards
     // Make sure that after input is propagated down, it doesn't have parents
     // that comes after i but before the new child
     int earliest_child = INT_MAX;
     const auto& iter = children.find(out);
     if (iter != children.end()) {
       for (int child : iter->second) {
         earliest_child = std::min(earliest_child, child);
       }
     }
     if (earliest_child == INT_MAX) {
       return true;
     }
     const auto& iter2 = parents.find(in);
     if (iter2 != parents.end()) {
       for (int parent : iter2->second) {
         if (parent > op_idx && parent < earliest_child) {
           VLOG(2) << "InplaceOps(" << op_type << ") skipping op: \n"
                   << op->DebugString();
           return false;
         }
       }
     }
   } else {
     // Propagate output upwards
     if (inputs.count(in) != 0 || outputs.count(in) != 0) {
       // This is the case when the op is absolutely needed. It exists to serve
       // one and only one purpose, to copy from in to out where in is one of
       // the net's inputs or outputs and out is one of the net's outputs.
       VLOG(2) << "InplaceOps(" << op_type << ") skipping op: \n"
               << op->DebugString();
       return false;
     }
     // find latest parent of in
     int latest_parent = -1;
     const auto& iter = parents.find(in);
     if (iter != parents.end()) {
       for (int parent : iter->second) {
         latest_parent = std::max(latest_parent, parent);
       }
     }
     if (latest_parent == -1) {
       return false;
     }
     // Make sure that after output is propagated, it doesn't have children that
     // comes after its new parent, but before its previous parent
     const auto& iter2 = children.find(out);
     if (iter2 != children.end()) {
       for (int child : iter2->second) {
         if (child < op_idx && child > latest_parent) {
           VLOG(2) << "InplaceOps(" << op_type << ") skipping op: \n"
                   << op->DebugString();
           return false;
         }
       }
     }
   }

   return true;
 }

 } // namespace

 // Conceptually it's a pretty easy process and consists of 3 steps:
 // 1) SSA rewrite; 2) propagate inputs forwards; 3) propagate outputs
 // backwards and then forwards again. However, because of model outputs
 // which can't be overwritten during the SSA process, and the fact that
 // inputs could be overwritten by ops and also appear in outputs, it adds
 // a lot of extra complexity to handle these special cases. A lot of this
 // extra complexity is handled in FoundOpCandidate.
 void RemoveOpsByType(InferenceGraph& graph, const std::string& op_type) {
   int num_removed = 0;
   NetDef* net = graph.predict_net_def.get();
   for (auto& op : net->op()) {
     if (op.type() == "RecurrentNetwork") {
       LOG(INFO) << "RemoveOpsByType does not support RecurrentNetwork yet";
       return;
     }
   }

   std::unordered_set<std::string> inputs(
       graph.input_names.begin(), graph.input_names.end());
   std::unordered_set<std::string> outputs(
       graph.output_names.begin(), graph.output_names.end());

   if (!graph.predictor_net_ssa_rewritten) {
     net->mutable_external_output()->Clear();
     // add external_outputs to net as they're necessary to correctly do ssa
     // rewriting
     for (const auto& o : graph.output_names) {
       net->add_external_output(o);
     }
     onnx::SsaRewrite(nullptr, net);
     // clear external_outputs
     net->mutable_external_output()->Clear();
     graph.predictor_net_ssa_rewritten = true;
   }

   // construct parents/children graphs to facilitate graph traversal
   std::unordered_map<std::string, std::unordered_set<int>> parents, children;
   for (int i = 0; i < net->op_size(); i++) {
     OperatorDef* op = net->mutable_op(i);
     for (auto& in : op->input()) {
       children[in].insert(i);
     }
     for (auto& output : op->output()) {
       parents[output].insert(i);
     }
   }

   // Inplace ops. Step 1: propagate inputs downward
   for (int i = 0; i < net->op_size(); i++) {
     OperatorDef* op = net->mutable_op(i);
     if (!FoundOpCandidate(op, i, op_type, inputs, outputs, parents, children)) {
       continue;
     }
     const std::string in = op->input(0);
     const std::string out = op->output(0);
     if (outputs.count(out) == 0) {
       // Rename all apperances of out to in
       VLOG(2) << "InplaceOps(" << op_type << ") inplacing op:\n"
               << op->DebugString();
       RenameInputsInChildren(out, in, net, children);
       RenameOutputs(out, in, op, i, parents);
     }
   }

   // Step 2: propagate outputs upward
   for (int i = 0; i < net->op_size(); i++) {
     OperatorDef* op = net->mutable_op(i);
     if (!FoundOpCandidate(op, i, op_type, inputs, outputs, parents, children)) {
       continue;
     }
     const std::string in = op->input(0);
     const std::string out = op->output(0);
     if (outputs.count(out) != 0) {
       if (inputs.count(in) == 0 && outputs.count(in) == 0) {
         // Rename all apperances (regardless of inputs/outputs) of in (if not
         // in inputs) to out, when out is guaranteed to be produced a parent
         // op. With the parents/children graph which remembers all apprerances
         // of nodes (not just immediate parent/children), we don't need to
         // propagate the outputs back down again because those cases are already
         // handled by RenameOutputInParents and RenameInputsInChildren
         if (parents.count(in) > 0 && !parents[in].empty()) {
           RenameOutputInParents(in, out, net, parents);
           VLOG(2) << "InplaceOps(" << op_type << ") inplacing op:\n"
                   << op->DebugString();
           RenameInputsInChildren(in, out, net, children);
           RenameInputs(in, out, op, i, children);
         }
       }
     }
   }

   // Remove inplace ops
   int i = 0;
   while (i < net->op_size()) {
     OperatorDef op = net->op(i);
     if (op.type() == op_type && op.input_size() == 1 && op.output_size() == 1 &&
         op.input(0) == op.output(0)) {
       net->mutable_op()->erase(net->mutable_op()->begin() + i);
       num_removed++;
       VLOG(2) << "RemoveOpsByType(" << op_type << ") deleting inplace op: \n"
               << op.DebugString();
     } else {
       i++;
       VLOG(2) << "RemoveOpsByType(" << op_type << ") skipping op: \n"
               << op.DebugString();
     }
   }
   VLOG(2) << "RemoveOpsByType(" << op_type << ") removed " << num_removed
           << " ops";
 }
 } // namespace caffe2
	#include "caffe2/predictor/transforms.h"
	#include "caffe2/onnx/onnx_exporter.h"
	#include "caffe2/utils/proto_utils.h"

	#include <unordered_set>

	namespace caffe2 {

	namespace {
	bool HasInput(const string& blob, const OperatorDef& op) {
	for (const auto& in : op.input()) {
	if (blob == in) {
	return true;
	}
	}
	return false;
	}

	bool HasOutput(const string& blob, const OperatorDef& op) {
	for (const auto& out : op.output()) {
	if (blob == out) {
	return true;
	}
	}
	return false;
	}

	void RewriteSubnetsForIfOp(
	const string& from,
	const string& to,
	OperatorDef* op) {
	ArgumentHelper helper(*op);
	Argument then_arg = nullptr, else_arg = nullptr;

	std::map<std::string, std::string> oldname_to_newname;
	oldname_to_newname[from] = to;

	if (helper.HasSingleArgumentOfType<NetDef>("then_net")) {
	then_arg = GetMutableArgument("then_net", false, op);
	onnx::rewriteSubnet(then_arg, oldname_to_newname);
	}
	if (helper.HasSingleArgumentOfType<NetDef>("else_net")) {
	else_arg = GetMutableArgument("else_net", false, op);
	onnx::rewriteSubnet(else_arg, oldname_to_newname);
	}
	}

	void RenameInputs(
	const string& from,
	const string& to,
	OperatorDef* def,
	int op_idx,
	std::unordered_map<std::string, std::unordered_set<int>>& children) {
	VLOG(2) << "RenameInputs (from=" << from << ", to=" << to << ", "
	<< def->DebugString() << ")";
	for (int i = 0; i < def->input_size(); i++) {
	if (def->input(i) == from) {
	*def->mutable_input(i) = to;
	children[from].erase(op_idx);
	children[to].insert(op_idx);
	}
	}
	// Rename inputs in the subnets of If/AsyncIf op
	if (def->type() == "If" \|\| def->type() == "AsyncIf") {
	RewriteSubnetsForIfOp(from, to, def);
	}
	}

	void RenameOutputs(
	const string& from,
	const string& to,
	OperatorDef* def,
	int op_idx,
	std::unordered_map<std::string, std::unordered_set<int>>& parents) {
	VLOG(2) << "RenameOutputs (from=" << from << ", to=" << to << ", "
	<< def->DebugString() << ")";
	for (string& output : *def->mutable_output()) {
	if (output == from) {
	output = to;
	parents[from].erase(op_idx);
	parents[to].insert(op_idx);
	}
	}
	// Rename outputs in the subnets of If/AsyncIf op
	if (def->type() == "If" \|\| def->type() == "AsyncIf") {
	RewriteSubnetsForIfOp(from, to, def);
	}
	}

	void RenameInputsInChildren(
	const string& from,
	const string& to,
	caffe2::NetDef* net,
	std::unordered_map<std::string, std::unordered_set<int>>& children) {
	VLOG(2) << "RenameInputsInChildren (from=" << from << ", to=" << to << ")";
	if (children.count(from) == 0) {
	return;
	}

	// make an temporary copy here because we're going to modify children
	for (int child : std::unordered_set<int>(children[from])) {
	RenameInputs(from, to, net->mutable_op(child), child, children);
	}
	}

	void RenameOutputInParents(
	const std::string& from,
	const std::string& to,
	caffe2::NetDef* net,
	std::unordered_map<std::string, std::unordered_set<int>>& parents) {
	VLOG(2) << "RenameOutputInParents (from=" << from << ", to=" << to << ")";
	if (parents.count(from) == 0) {
	return;
	}

	// make an temporary copy here because we're going to modify parents
	for (int parent : std::unordered_set<int>(parents[from])) {
	RenameOutputs(from, to, net->mutable_op(parent), parent, parents);
	}
	}

	bool FoundOpCandidate(
	const OperatorDef* op,
	int op_idx,
	const std::string& op_type,
	const std::unordered_set<std::string>& inputs,
	const std::unordered_set<std::string>& outputs,
	const std::unordered_map<std::string, std::unordered_set<int>>& parents,
	const std::unordered_map<std::string, std::unordered_set<int>>& children) {
	if (op->type() != op_type) {
	VLOG(2) << "InplaceOps(" << op_type << ") skipping op: \n"
	<< op->DebugString();
	return false;
	}
	if (op->input_size() != 1 \|\| op->output_size() != 1) {
	VLOG(2) << "InplaceOps(" << op_type
	<< ") only supports ops with exactly 1 output "
	<< "and exactly 1 input. Skipping op: \n"
	<< op->DebugString();
	return false;
	}

	// use actual copy because op->input/output may change
	const std::string in = op->input(0);
	const std::string out = op->output(0);

	if (in == out) {
	// This case can still exist when in/out is in the predict_net's outputs.
	// The op is an inplace op already.
	return false;
	}

	// The following is to handle the special cases of inputs being overwritten
	// by ops in the net and then appear in outputs of the net
	if (outputs.count(out) == 0) {
	// Propagate input downwards
	// Make sure that after input is propagated down, it doesn't have parents
	// that comes after i but before the new child
	int earliest_child = INT_MAX;
	const auto& iter = children.find(out);
	if (iter != children.end()) {
	for (int child : iter->second) {
	earliest_child = std::min(earliest_child, child);
	}
	}
	if (earliest_child == INT_MAX) {
	return true;
	}
	const auto& iter2 = parents.find(in);
	if (iter2 != parents.end()) {
	for (int parent : iter2->second) {
	if (parent > op_idx && parent < earliest_child) {
	VLOG(2) << "InplaceOps(" << op_type << ") skipping op: \n"
	<< op->DebugString();
	return false;
	}
	}
	}
	} else {
	// Propagate output upwards
	if (inputs.count(in) != 0 \|\| outputs.count(in) != 0) {
	// This is the case when the op is absolutely needed. It exists to serve
	// one and only one purpose, to copy from in to out where in is one of
	// the net's inputs or outputs and out is one of the net's outputs.
	VLOG(2) << "InplaceOps(" << op_type << ") skipping op: \n"
	<< op->DebugString();
	return false;
	}
	// find latest parent of in
	int latest_parent = -1;
	const auto& iter = parents.find(in);
	if (iter != parents.end()) {
	for (int parent : iter->second) {
	latest_parent = std::max(latest_parent, parent);
	}
	}
	if (latest_parent == -1) {
	return false;
	}
	// Make sure that after output is propagated, it doesn't have children that
	// comes after its new parent, but before its previous parent
	const auto& iter2 = children.find(out);
	if (iter2 != children.end()) {
	for (int child : iter2->second) {
	if (child < op_idx && child > latest_parent) {
	VLOG(2) << "InplaceOps(" << op_type << ") skipping op: \n"
	<< op->DebugString();
	return false;
	}
	}
	}
	}

	return true;
	}

	} // namespace

	// Conceptually it's a pretty easy process and consists of 3 steps:
	// 1) SSA rewrite; 2) propagate inputs forwards; 3) propagate outputs
	// backwards and then forwards again. However, because of model outputs
	// which can't be overwritten during the SSA process, and the fact that
	// inputs could be overwritten by ops and also appear in outputs, it adds
	// a lot of extra complexity to handle these special cases. A lot of this
	// extra complexity is handled in FoundOpCandidate.
	void RemoveOpsByType(InferenceGraph& graph, const std::string& op_type) {
	int num_removed = 0;
	NetDef* net = graph.predict_net_def.get();
	for (auto& op : net->op()) {
	if (op.type() == "RecurrentNetwork") {
	LOG(INFO) << "RemoveOpsByType does not support RecurrentNetwork yet";
	return;
	}
	}

	std::unordered_set<std::string> inputs(
	graph.input_names.begin(), graph.input_names.end());
	std::unordered_set<std::string> outputs(
	graph.output_names.begin(), graph.output_names.end());

	if (!graph.predictor_net_ssa_rewritten) {
	net->mutable_external_output()->Clear();
	// add external_outputs to net as they're necessary to correctly do ssa
	// rewriting
	for (const auto& o : graph.output_names) {
	net->add_external_output(o);
	}
	onnx::SsaRewrite(nullptr, net);
	// clear external_outputs
	net->mutable_external_output()->Clear();
	graph.predictor_net_ssa_rewritten = true;
	}

	// construct parents/children graphs to facilitate graph traversal
	std::unordered_map<std::string, std::unordered_set<int>> parents, children;
	for (int i = 0; i < net->op_size(); i++) {
	OperatorDef* op = net->mutable_op(i);
	for (auto& in : op->input()) {
	children[in].insert(i);
	}
	for (auto& output : op->output()) {
	parents[output].insert(i);
	}
	}

	// Inplace ops. Step 1: propagate inputs downward
	for (int i = 0; i < net->op_size(); i++) {
	OperatorDef* op = net->mutable_op(i);
	if (!FoundOpCandidate(op, i, op_type, inputs, outputs, parents, children)) {
	continue;
	}
	const std::string in = op->input(0);
	const std::string out = op->output(0);
	if (outputs.count(out) == 0) {
	// Rename all apperances of out to in
	VLOG(2) << "InplaceOps(" << op_type << ") inplacing op:\n"
	<< op->DebugString();
	RenameInputsInChildren(out, in, net, children);
	RenameOutputs(out, in, op, i, parents);
	}
	}

	// Step 2: propagate outputs upward
	for (int i = 0; i < net->op_size(); i++) {
	OperatorDef* op = net->mutable_op(i);
	if (!FoundOpCandidate(op, i, op_type, inputs, outputs, parents, children)) {
	continue;
	}
	const std::string in = op->input(0);
	const std::string out = op->output(0);
	if (outputs.count(out) != 0) {
	if (inputs.count(in) == 0 && outputs.count(in) == 0) {
	// Rename all apperances (regardless of inputs/outputs) of in (if not
	// in inputs) to out, when out is guaranteed to be produced a parent
	// op. With the parents/children graph which remembers all apprerances
	// of nodes (not just immediate parent/children), we don't need to
	// propagate the outputs back down again because those cases are already
	// handled by RenameOutputInParents and RenameInputsInChildren
	if (parents.count(in) > 0 && !parents[in].empty()) {
	RenameOutputInParents(in, out, net, parents);
	VLOG(2) << "InplaceOps(" << op_type << ") inplacing op:\n"
	<< op->DebugString();
	RenameInputsInChildren(in, out, net, children);
	RenameInputs(in, out, op, i, children);
	}
	}
	}
	}

	// Remove inplace ops
	int i = 0;
	while (i < net->op_size()) {
	OperatorDef op = net->op(i);
	if (op.type() == op_type && op.input_size() == 1 && op.output_size() == 1 &&
	op.input(0) == op.output(0)) {
	net->mutable_op()->erase(net->mutable_op()->begin() + i);
	num_removed++;
	VLOG(2) << "RemoveOpsByType(" << op_type << ") deleting inplace op: \n"
	<< op.DebugString();
	} else {
	i++;
	VLOG(2) << "RemoveOpsByType(" << op_type << ") skipping op: \n"
	<< op.DebugString();
	}
	}
	VLOG(2) << "RemoveOpsByType(" << op_type << ") removed " << num_removed
	<< " ops";
	}
	} // namespace caffe2