mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
0f419abf40
Summary: We need to remove nomnigraph from the list of public libraries in order to support libtorch extensions. Easiest way to do this is to include it into the Caffe2 source like all other caffe2/core/ code. However, because the headers are in a different place, we need to include them for linked libraries (pybind, tests, etc). On an upside, this means that nomnigraph is now default hidden visibility too. FYI peterjc123 xkszltl goldsborough bwasti Yangqing Pull Request resolved: https://github.com/pytorch/pytorch/pull/11303 Reviewed By: pjh5 Differential Revision: D9694932 Pulled By: orionr fbshipit-source-id: 5db3eb20bc5ddc873ce9151236b74663fbb33ed8
38 lines
1.3 KiB
C++
38 lines
1.3 KiB
C++
#include "caffe2/core/common.h"
|
|
|
|
class CAFFE2_API OnnxAnnotation : public nom::repr::Annotation {
|
|
public:
|
|
OnnxAnnotation() : Annotation(AnnotationKind::Onnx) {}
|
|
OnnxAnnotation(std::string device)
|
|
: Annotation(AnnotationKind::Onnx), Device(device) {}
|
|
|
|
void setDevice(std::string device) { Device = device; }
|
|
const std::string getDevice() const { return Device; }
|
|
|
|
void setOperatorDef(caffe2::OperatorDef* opDef) {
|
|
OpDef = opDef;
|
|
}
|
|
const caffe2::OperatorDef* getOperatorDef() const {
|
|
assert(OpDef && "OperatorDef was never set. Use OnnxAnnotation::setOperatorDef.");
|
|
return OpDef;
|
|
}
|
|
caffe2::OperatorDef* getMutableOperatorDef() {
|
|
assert(OpDef && "OperatorDef was never set. Use OnnxAnnotation::setOperatorDef.");
|
|
return OpDef;
|
|
}
|
|
|
|
static bool classof(const Annotation *A) {
|
|
return A->getKind() == AnnotationKind::Onnx;
|
|
}
|
|
|
|
private:
|
|
std::string Device = "";
|
|
caffe2::OperatorDef* OpDef = nullptr;
|
|
};
|
|
|
|
CAFFE2_API nom::repr::NNModule convertToNNModule(caffe2::NetDef &net, std::unordered_map<std::string, nom::repr::NNGraph::NodeRef>* blobMapOut = nullptr);
|
|
|
|
CAFFE2_API caffe2::NetDef convertToOnnxProto(nom::repr::NNModule&);
|
|
|
|
CAFFE2_API std::unique_ptr<nom::repr::NeuralNetOperator> convertToOperatorDef(caffe2::OperatorDef op);
|