#include <test/cpp/jit/test_alias_analysis.h>
#include <test/cpp/jit/test_misc.h>
-#include <test/cpp/jit/test_netdef_converter.h>
using namespace torch;
using namespace torch::jit;
JIT_TEST(AliasAnalysis)
JIT_TEST(AliasTracker)
-JIT_TEST(NetDefConverter)
-
JIT_TEST(THNNConv)
JIT_TEST(ATenNativeBatchNorm)
#include <test/cpp/jit/test_alias_analysis.h>
#include <test/cpp/jit/test_misc.h>
-#include <test/cpp/jit/test_netdef_converter.h>
#include <sstream>
#include <string>
testRegisterFusionCachesKernel();
testAliasAnalysis();
testAliasTracker();
- testNetDefConverter(out);
return out.str();
}
+++ /dev/null
-#pragma once
-
-#include <torch/csrc/jit/netdef_converter.h>
-#include "test/cpp/jit/test_base.h"
-
-#include <sstream>
-#include <string>
-
-namespace torch {
-namespace jit {
-
-void testNetDefConverter(std::ostream& out = std::cout) {
- {
- // Check a simple net conversion back and forth.
-
- // Create a simple graph:
- // graph(%0 : Tensor
- // %1 : Tensor) {
- // %2 : Tensor = aten::mul(%0, %1)
- // %3 : int = prim::Constant[value=1]()
- // %4 : Tensor = aten::add(%0, %2, %3)
- // return (%2, %4);
- // }
- auto graph = std::make_shared<Graph>();
- auto a = graph->addInput();
- auto b = graph->addInput();
- auto c = graph->insert(aten::mul, {a, b});
- auto d = graph->insert(aten::add, {a, c});
- graph->registerOutput(c);
- graph->registerOutput(d);
-
- // Convert it to netdef and check the result
- caffe2::NetDef net;
- convertIRToNetDef(&net, *graph);
- AT_ASSERT(net.op().size() == 3);
- AT_ASSERT(net.external_input().size() == 2);
- AT_ASSERT(net.external_output().size() == 2);
-
- const caffe2::OperatorDef& MulOp = net.op().Get(0);
- AT_ASSERT(MulOp.input().size() == 2);
- AT_ASSERT(MulOp.input().Get(0) == net.external_input().Get(0));
- AT_ASSERT(MulOp.input().Get(1) == net.external_input().Get(1));
- AT_ASSERT(MulOp.output().size() == 1);
-
- const caffe2::OperatorDef& ConstNode = net.op().Get(1);
- AT_ASSERT(ConstNode.input().size() == 0);
- AT_ASSERT(ConstNode.output().size() == 1);
- AT_ASSERT(ConstNode.arg().size() == 1);
- AT_ASSERT(ConstNode.arg().Get(0).name() == "value");
- AT_ASSERT(ConstNode.arg().Get(0).i() == 1);
-
- const caffe2::OperatorDef& AddOp = net.op().Get(2);
- AT_ASSERT(AddOp.input().size() == 3);
- AT_ASSERT(AddOp.input().Get(0) == net.external_input().Get(0));
- AT_ASSERT(AddOp.input().Get(1) == MulOp.output().Get(0));
- AT_ASSERT(AddOp.input().Get(2) == ConstNode.output().Get(0));
-
- AT_ASSERT(net.external_output().Get(0) == MulOp.output().Get(0));
- AT_ASSERT(net.external_output().Get(1) == AddOp.output().Get(0));
-
- // Convert NetDef back to IR and check if we get the original.
- Graph graph2;
- std::unordered_map<std::string, Value*> vmap;
- convertNetDefToIR(net, &graph2, &vmap);
-
- Node* mul = graph2.outputs()[0]->node();
- Node* add = graph2.outputs()[1]->node();
- AT_ASSERT(mul->kind() == c->node()->kind());
- AT_ASSERT(add->kind() == d->node()->kind());
- AT_ASSERT(mul->inputs()[0] == graph2.inputs()[0]);
- AT_ASSERT(mul->inputs()[1] == graph2.inputs()[1]);
- AT_ASSERT(add->inputs()[0] == graph2.inputs()[0]);
- AT_ASSERT(add->inputs()[1] == graph2.outputs()[0]);
- }
- {
- // Check attributes conversion
- auto graph = std::make_shared<Graph>();
- auto a = graph->addInput();
- auto b = graph->addInput();
- Node* node =
- graph->create(Symbol::fromQualString("test::some_op"), {a, b}, 2);
- graph->insertNode(node);
-
- node->i_(Symbol::fromQualString("attr::i_attr"), 42);
- node->f_(Symbol::fromQualString("attr::f_attr"), 3.0);
- node->s_(Symbol::fromQualString("attr::s_attr"), "Hello!");
-
- node->is_(Symbol::fromQualString("attr::is_attr"), {14, 18, 7});
- node->fs_(Symbol::fromQualString("attr::fs_attr"), {2.72, 3.14});
- node->ss_(Symbol::fromQualString("attr::ss_attr"), {"Winter", "Summer"});
-
- graph->registerOutput(node->outputs()[0]);
- graph->registerOutput(node->outputs()[1]);
-
- // Convert it to netdef and check the result
- caffe2::NetDef net;
- convertIRToNetDef(&net, *graph);
- const caffe2::OperatorDef& Op = net.op().Get(0);
- AT_ASSERT(Op.arg().Get(0).name() == "i_attr");
- AT_ASSERT(Op.arg().Get(0).i() == 42);
- AT_ASSERT(Op.arg().Get(1).name() == "f_attr");
- AT_ASSERT(Op.arg().Get(1).f() == 3.0);
- AT_ASSERT(Op.arg().Get(2).name() == "s_attr");
- AT_ASSERT(Op.arg().Get(2).s() == "Hello!");
-
- AT_ASSERT(Op.arg().Get(3).name() == "is_attr");
- AT_ASSERT(Op.arg().Get(3).ints().size() == 3);
- AT_ASSERT(Op.arg().Get(3).ints().Get(0) == 14);
- AT_ASSERT(Op.arg().Get(3).ints().Get(1) == 18);
- AT_ASSERT(Op.arg().Get(3).ints().Get(2) == 7);
-
- AT_ASSERT(Op.arg().Get(4).name() == "fs_attr");
- AT_ASSERT(Op.arg().Get(4).floats().size() == 2);
- AT_ASSERT(fabs(Op.arg().Get(4).floats().Get(0) - 2.72) < 0.001);
-
- AT_ASSERT(Op.arg().Get(5).name() == "ss_attr");
- AT_ASSERT(Op.arg().Get(5).strings().size() == 2);
- AT_ASSERT(Op.arg().Get(5).strings().Get(1) == "Summer");
-
- AT_ASSERT(net.external_output().Get(0) == Op.output().Get(0));
- AT_ASSERT(net.external_output().Get(1) == Op.output().Get(1));
-
- // Convert NetDef back to IR and check if we get the original.
- Graph graph2;
- std::unordered_map<std::string, Value*> vmap;
- convertNetDefToIR(net, &graph2, &vmap);
-
- AT_ASSERT(graph2.outputs()[0]->node() == graph2.outputs()[0]->node());
- Node* n = graph2.outputs()[0]->node();
- AT_ASSERT(n->i(Symbol::fromQualString("attr::i_attr")) == 42);
- AT_ASSERT(n->f(Symbol::fromQualString("attr::f_attr")) == 3.0);
- AT_ASSERT(n->s(Symbol::fromQualString("attr::s_attr")) == "Hello!");
- AT_ASSERT(
- n->is(Symbol::fromQualString("attr::is_attr")) ==
- std::vector<long>({14, 18, 7}));
- AT_ASSERT(
- fabs(n->fs(Symbol::fromQualString("attr::fs_attr"))[0] - 2.72) < 0.001);
- AT_ASSERT(
- fabs(n->fs(Symbol::fromQualString("attr::fs_attr"))[1] - 3.14) < 0.001);
- AT_ASSERT(
- n->ss(Symbol::fromQualString("attr::ss_attr")) ==
- std::vector<std::string>({"Winter", "Summer"}));
- }
-}
-} // namespace jit
-} // namespace torch
${TORCH_SRC_DIR}/csrc/jit/constants.cpp
${TORCH_SRC_DIR}/csrc/jit/node_hashing.cpp
${TORCH_SRC_DIR}/csrc/jit/ir.cpp
- ${TORCH_SRC_DIR}/csrc/jit/netdef_converter.cpp
${TORCH_SRC_DIR}/csrc/jit/operator.cpp
${TORCH_SRC_DIR}/csrc/jit/caffe2_operator.cpp
${TORCH_SRC_DIR}/csrc/jit/register_c10_ops.cpp
+++ /dev/null
-#include <torch/csrc/jit/netdef_converter.h>
-
-namespace torch {
-namespace jit {
-
-static AttributeKind getArgKind(const caffe2::Argument& arg) {
- if (arg.has_i()) {
- return AttributeKind::i;
- } else if (arg.has_f()) {
- return AttributeKind::f;
- } else if (arg.has_s()) {
- return AttributeKind::s;
- } else if (arg.has_t()) {
- return AttributeKind::t;
- } else if (arg.has_n()) {
- return AttributeKind::g;
- } else if (arg.ints().size()) {
- return AttributeKind::is;
- } else if (arg.floats().size()) {
- return AttributeKind::fs;
- } else if (arg.strings().size()) {
- return AttributeKind::ss;
- } else if (arg.tensors().size()) {
- return AttributeKind::ts;
- } else if (arg.nets().size()) {
- return AttributeKind::gs;
- }
- // Unknown type.
- abort();
-}
-
-static void convertArg(const caffe2::Argument& arg, Node* node) {
- std::string attrName = "attr::" + arg.name();
- auto attrSymbol = Symbol::fromQualString(attrName);
- AttributeKind kind = getArgKind(arg);
- switch (kind) {
- case AttributeKind::i: {
- node->i_(attrSymbol, (long)arg.i());
- break;
- }
- case AttributeKind::f: {
- node->f_(attrSymbol, arg.f());
- break;
- }
- case AttributeKind::s: {
- node->s_(attrSymbol, arg.s());
- break;
- }
- case AttributeKind::is: {
- std::vector<long> is(arg.ints().begin(), arg.ints().end());
- node->is_(attrSymbol, is);
- break;
- }
- case AttributeKind::fs: {
- std::vector<double> fs(arg.floats().begin(), arg.floats().end());
- node->fs_(attrSymbol, fs);
- break;
- }
- case AttributeKind::ss: {
- std::vector<std::string> ss(arg.strings().begin(), arg.strings().end());
- node->ss_(attrSymbol, ss);
- break;
- }
- default: {
- std::cout << "Unsupported type '" << toString(kind) << "' of attribute '"
- << attrName << "'"
- << " in node:" << std::endl;
- node->dump();
- abort();
- }
- }
-}
-
-void convertNetDefToIR(
- const caffe2::NetDef& net,
- Graph* g,
- std::unordered_map<std::string, Value*>* valueMapPtr,
- const std::string& prefix) {
- std::unordered_map<std::string, Value*>& valueMap = *valueMapPtr;
- valueMap.clear();
-
- for (const auto& inputName : net.external_input()) {
- AT_ASSERT(!valueMap.count(inputName));
- valueMap[inputName] = g->addInput();
- }
-
- for (const auto& op : net.op()) {
- std::string name = prefix + op.type();
- Node* node =
- g->create(Symbol::fromQualString(name), {}, op.output().size());
- g->insertNode(node);
-
- for (const auto& input : op.input()) {
- AT_ASSERT(valueMap.count(input));
- node->addInput(valueMap[input]);
- }
- int idx = 0;
- for (const auto& output : op.output()) {
- // If output already exists in valueMap, overwrite it. This way we will
- // have the last definition of a value named 'output' in valueMap.
- valueMap[output] = node->outputs()[idx++];
- }
- for (const auto& arg : op.arg()) {
- convertArg(arg, node);
- }
- }
-
- for (const auto& outputName : net.external_output()) {
- AT_ASSERT(valueMap.count(outputName));
- g->registerOutput(valueMap.at(outputName));
- }
-}
-
-static void convertAttrToCaffe2Arg(
- const Node* node,
- const Symbol& name,
- caffe2::Argument* arg) {
- arg->set_name(name.toUnqualString());
- switch (node->kindOf(name)) {
- case AttributeKind::i: {
- arg->set_i(node->i(name));
- break;
- }
- case AttributeKind::f: {
- arg->set_f(node->f(name));
- break;
- }
- case AttributeKind::s: {
- arg->set_s(node->s(name));
- break;
- }
- case AttributeKind::is: {
- for (long i : node->is(name)) {
- arg->add_ints(i);
- }
- break;
- }
- case AttributeKind::fs: {
- for (double f : node->fs(name)) {
- arg->add_floats(f);
- }
- break;
- }
- case AttributeKind::ss: {
- for (const std::string& s : node->ss(name)) {
- arg->add_strings(s);
- }
- break;
- }
- default: {
- std::cout << "Unsupported type '" << toString(node->kindOf(name))
- << "' of attribute '" << name.toUnqualString() << "'"
- << " in node:" << std::endl;
- node->dump();
- abort();
- }
- }
-}
-
-static void convertNodeToCaffe2Op(const Node* node, caffe2::NetDef* net) {
- caffe2::OperatorDef op;
- op.set_type(node->kind().toQualString());
- for (const Value* input : node->inputs()) {
- op.add_input(input->uniqueName());
- }
- for (const Value* output : node->outputs()) {
- op.add_output(output->uniqueName());
- }
- std::vector<Symbol> names = node->attributeNames();
- for (const Symbol& name : names) {
- caffe2::Argument* arg = op.add_arg();
- convertAttrToCaffe2Arg(node, name, arg);
- }
- *net->add_op() = op;
-}
-
-void convertIRToNetDef(caffe2::NetDef* net, const Graph& g) {
- net->mutable_op()->Clear();
-
- for (const Value* value : g.inputs()) {
- net->add_external_input(value->uniqueName());
- }
-
- for (const Node* node : g.nodes()) {
- convertNodeToCaffe2Op(node, net);
- }
-
- for (const Value* value : g.outputs()) {
- net->add_external_output(value->uniqueName());
- }
-}
-
-} // namespace jit
-} // namespace torch
+++ /dev/null
-#pragma once
-#include <caffe2/proto/caffe2_pb.h>
-#include <torch/csrc/jit/ir.h>
-#include <unordered_map>
-
-namespace torch {
-namespace jit {
-
-/** \brief Convert a caffe2 NetDef to PyTorch IR.
- *
- * The NetDef \p net is converted and the result is stored in the
- * torch::jit::Graph \p graph. The function also records name->value map in \p
- * valueMapPtr. If the original net had several values with the same name, the
- * map will contain the value for the last definition.
- * \p Prefix can be used for appending some string to every operator name (e.g.
- * we can add "caffe2::").
- */
-void convertNetDefToIR(
- const caffe2::NetDef& net,
- Graph* graph,
- std::unordered_map<std::string, Value*>* valueMapPtr,
- const std::string& prefix = "");
-
-/** \brief Convert PyTorch IR \p graph to Caffe2 NetDef \p net.
- *
- * Note: for constant nodes (prim::Const) we generate a separate op in the net,
- * which might or might not be what we want. The idea here is that eventually
- * both formats will converge to PyTorch IR, so for now we try to keep as close
- * to it as possible. For short-term applications we might add a separate pass
- * that would fold such const-nodes into their users.
- *
- * TODO: We might need to do a better job at preserving names of the variables,
- * especially external_inputs/external_outputs.
- */
-void convertIRToNetDef(caffe2::NetDef* net, const Graph& graph);
-
-} // namespace jit
-} // namespace torch
projects / platform / upstream / pytorch.git / commitdiff