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
69 lines
2.1 KiB
C++
69 lines
2.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 {
|
|
|
|
class TORCH_API Caffe2Annotation : public nom::repr::Annotation {
|
|
public:
|
|
Caffe2Annotation() : Annotation(AnnotationKind::Caffe2) {}
|
|
Caffe2Annotation(std::string device)
|
|
: Annotation(AnnotationKind::Caffe2), Device(device) {}
|
|
virtual ~Caffe2Annotation() {}
|
|
|
|
void setOperatorDef(const caffe2::OperatorDef& opDef);
|
|
bool hasOperatorDef() const;
|
|
const caffe2::OperatorDef& getOperatorDef() const;
|
|
caffe2::OperatorDef* getMutableOperatorDef();
|
|
|
|
void setDeviceOption(const caffe2::DeviceOption& opDef);
|
|
bool hasDeviceOption() const;
|
|
const caffe2::DeviceOption& getDeviceOption() const;
|
|
caffe2::DeviceOption* getMutableDeviceOption();
|
|
|
|
// Distributed annotations
|
|
void setDevice(std::string device);
|
|
const std::string getDevice() const;
|
|
void setDeviceType(int device);
|
|
int getDeviceType() const;
|
|
|
|
enum class ParallelizationScheme {
|
|
none,
|
|
split_by_batch,
|
|
split_by_length,
|
|
shard,
|
|
shard_by_number
|
|
};
|
|
void setParallelization(ParallelizationScheme, int num = -1);
|
|
ParallelizationScheme getParallelizationScheme() const;
|
|
int getParallelization() const;
|
|
|
|
void setKeyNode(nom::repr::NNGraph::NodeRef);
|
|
const nom::repr::NNGraph::NodeRef& getKeyNode() const;
|
|
void setLengthNode(nom::repr::NNGraph::NodeRef);
|
|
const nom::repr::NNGraph::NodeRef& getLengthNode() const;
|
|
|
|
void setComponentLevels(std::vector<std::string> components);
|
|
std::vector<std::string> getComponentLevels() const;
|
|
|
|
static bool classof(const Annotation* A);
|
|
|
|
private:
|
|
std::string Device = "";
|
|
caffe2::OperatorDef OpDef;
|
|
bool OpDefExists = false;
|
|
|
|
// Distributed annotations
|
|
int DeviceType = caffe2::DeviceTypeProto::PROTO_CPU;
|
|
ParallelizationScheme parallelization_scheme_ = ParallelizationScheme::none;
|
|
int parallelization_ = -1;
|
|
nom::repr::NNGraph::NodeRef key_node_ = nullptr;
|
|
nom::repr::NNGraph::NodeRef length_node_ = nullptr;
|
|
std::vector<std::string> component_levels_;
|
|
};
|
|
|
|
} // namespace caffe2
|