mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
8f0603b128
Summary: * adds TORCH_API and AT_CUDA_API in places * refactor code generation Python logic to separate caffe2/torch outputs * fix hip and asan * remove profiler_cuda from hip * fix gcc warnings for enums * Fix PythonOp::Kind Pull Request resolved: https://github.com/pytorch/pytorch/pull/19554 Differential Revision: D15082727 Pulled By: kostmo fbshipit-source-id: 83a8a99717f025ab44b29608848928d76b3147a4
45 lines
1.7 KiB
C++
45 lines
1.7 KiB
C++
#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
|