mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
71ca600af9
Summary: Since caffe2 and torch have been consolidated, CAFFE2_API should be merged with TORCH_API. Addresses a TODO. Manually edited some references of the removed `CAFFE2_API`: * `CONTRIBUTING.md` * `caffe2/proto/CMakeLists.txt` * `cmake/ProtoBuf.cmake` * `c10/macros/Export.h` * `torch/csrc/WindowsTorchApiMacro.h` Pull Request resolved: https://github.com/pytorch/pytorch/pull/49496 Reviewed By: malfet, samestep Differential Revision: D25600726 Pulled By: janeyx99 fbshipit-source-id: 7e068d959e397ac183c097d7e9a9afeca5ddd782
51 lines
1.1 KiB
C++
51 lines
1.1 KiB
C++
#pragma once
|
|
|
|
#include "caffe2/predictor/predictor.h"
|
|
#include "caffe2/proto/caffe2_pb.h"
|
|
|
|
#include <memory>
|
|
#include <string>
|
|
#include <vector>
|
|
|
|
namespace caffe2 {
|
|
namespace onnx {
|
|
class TORCH_API Caffe2BackendRep {
|
|
public:
|
|
void Run(
|
|
const caffe2::Predictor::TensorList& inputs,
|
|
caffe2::Predictor::TensorList* outputs);
|
|
void RunMap(
|
|
const caffe2::Predictor::TensorMap& inputs,
|
|
caffe2::Predictor::TensorList* outputs);
|
|
|
|
caffe2::NetDef& init_net() {
|
|
return init_net_;
|
|
}
|
|
caffe2::NetDef& pred_net() {
|
|
return pred_net_;
|
|
}
|
|
std::vector<std::string>& uninitialized_inputs() {
|
|
return uninitialized_inputs_;
|
|
}
|
|
|
|
const caffe2::NetDef& init_net() const {
|
|
return init_net_;
|
|
}
|
|
const caffe2::NetDef& pred_net() const {
|
|
return pred_net_;
|
|
}
|
|
const std::vector<std::string>& uninitialized_inputs() const {
|
|
return uninitialized_inputs_;
|
|
}
|
|
|
|
private:
|
|
void CheckInit();
|
|
|
|
caffe2::NetDef init_net_;
|
|
caffe2::NetDef pred_net_;
|
|
std::vector<std::string> uninitialized_inputs_;
|
|
std::unique_ptr<caffe2::Predictor> predictor_{nullptr};
|
|
};
|
|
} // namespace onnx
|
|
} // namespace caffe2
|