mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +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
34 lines
1.1 KiB
C++
34 lines
1.1 KiB
C++
#pragma once
|
|
|
|
#include "caffe2/core/common.h"
|
|
#include "caffe2/core/logging.h"
|
|
#include "caffe2/proto/caffe2_pb.h"
|
|
#include "nomnigraph/Representations/NeuralNet.h"
|
|
|
|
namespace caffe2 {
|
|
|
|
/// \brief Convert to an NNModule and apply a mapping of
|
|
/// tensor names to DeviceOptions to it.
|
|
///
|
|
/// This *only* applies the map to Declare/Export
|
|
/// nodes, which are representationally equivalent to
|
|
/// external_input/external_output in caffe2 NetDefs.
|
|
///
|
|
/// Throws an exception if the passed in blobMap contains
|
|
/// blobs that are not present in the NNModule.
|
|
TORCH_API nom::repr::NNModule convertToNNModule(
|
|
caffe2::NetDef&,
|
|
std::map<std::string, caffe2::DeviceOption>);
|
|
|
|
/// Helpers for the convertToNNModule for use
|
|
/// if you already have an NNModule.
|
|
/// You probably don't want to use these
|
|
/// if you can use convertToNNModule instead.
|
|
TORCH_API void addBlobDeviceOptions(
|
|
std::map<std::string, caffe2::DeviceOption> blobMap,
|
|
nom::repr::NNModule* nn);
|
|
TORCH_API void injectDataEdgeIndicators(nom::repr::NNModule* nn);
|
|
TORCH_API void removeDataEdgeIndicators(nom::repr::NNModule* nn);
|
|
|
|
} // namespace caffe2
|