1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
|
#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
|
