#pragma once
|
#include <caffe2/proto/caffe2_pb.h>
|
#include <torch/csrc/jit/ir.h>
|
#include <unordered_map>
|
#include <torch/csrc/WindowsTorchApiMacro.h>
|
|
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. valueMapPtr is optional.
|
* \p Prefix can be used for appending some string to every operator name (e.g.
|
* we can add "caffe2::").
|
*/
|
TORCH_API void convertNetDefToIR(
|
const caffe2::NetDef& net,
|
Graph* graph,
|
std::unordered_map<std::string, Value*>* valueMapPtr = nullptr,
|
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.
|
* \p If Prefix is specified, the prefix will be removed from operator name when
|
* converting from IR to NetDef.
|
*
|
* TODO: We might need to do a better job at preserving names of the variables,
|
* especially external_inputs/external_outputs.
|
*/
|
TORCH_API void convertIRToNetDef(
|
caffe2::NetDef* net,
|
const Graph& graph,
|
const std::string& prefix = "");
|
|
} // namespace jit
|
} // namespace torch
|
