mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
4a870f6518
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/49385 Currently, the API to export operator lists accepts a `torch::jit::Module` object, and spits out an operator list. The operator list is practically used only for mobile. This is not ideal because the set of root operators may change by the time the model is subsequently optmized and exported for mobile. What we need to to instead is glean the list of operators from the mobile model itself (`bytecode.pkl` specifically), and expose that instead. Also updated the logic in `converter`. ### Before this change: 1. Get operator List from Torch Script Model 2. Convert to bytecode mobile model ### After this change: 1. Convert to bytecode mobile model 2. Use this converted mobile model to get the list of operators for each method on the model ghstack-source-id: 118796752 Test Plan: Added a unit test in `test_lite_interpreter.cpp` to ensure that all model referenced operators show up in the exported operator list. Also make `test_lite_interpreter.cpp` runnable from `xplat/caffe2/BUCK` since this is where the production code will be built from. Verified that the list of operators produced before and after this change for an example model (segmentation) are the same. {P147863234} Also verified that the operator lists for BI-Xray model is different (we have been having problems with missing operators for this one): {P154903132} Reviewed By: iseeyuan Differential Revision: D24690094 fbshipit-source-id: 0426a6ef90456a811010cfe337c415882ae2deff
45 lines
1.2 KiB
C++
45 lines
1.2 KiB
C++
#pragma once
|
|
#include <ATen/core/ivalue.h>
|
|
//#include <aten/src/Aten/core/operator_name.h>
|
|
#include <vector>
|
|
|
|
namespace torch {
|
|
namespace jit {
|
|
using Stack = std::vector<c10::IValue>;
|
|
enum OpCode : uint8_t;
|
|
|
|
namespace mobile {
|
|
struct Code;
|
|
|
|
class Function {
|
|
public:
|
|
Function(c10::QualifiedName name);
|
|
bool run(Stack& stack) const;
|
|
c10::IValue operator()(Stack& stack);
|
|
const std::string& name() const;
|
|
const c10::QualifiedName& qualname() const;
|
|
void append_instruction(OpCode op, int X, int N);
|
|
bool append_operator(
|
|
const std::string& name,
|
|
const std::string& overload_name,
|
|
int64_t model_version);
|
|
void set_module_debug_info_list_size(size_t size);
|
|
void set_module_info(const std::string& module_info, size_t pc);
|
|
void append_constant(const c10::IValue& constant);
|
|
void append_type(const c10::TypePtr& type);
|
|
|
|
void set_register_size(size_t size);
|
|
|
|
std::string get_module_debug_info(size_t pc) const;
|
|
const std::shared_ptr<Code> get_code() const;
|
|
|
|
private:
|
|
c10::QualifiedName name_;
|
|
std::shared_ptr<Code> code_;
|
|
std::vector<std::string> pc_to_module_debug_info_;
|
|
};
|
|
|
|
} // namespace mobile
|
|
} // namespace jit
|
|
} // namespace torch
|