Skip to content

Commit 6b9e7b0

Browse files
David Linfacebook-github-bot
authored andcommitted
Add benchmark for training training (#6933)
Summary: This adds benchmarking for training. It loads the 3 methods required for training (param names, param index, gradient index) and runs the exported forward+backwards, then runs the optimizer. Reviewed By: JacobSzwejbka Differential Revision: D65905482
1 parent 809a1a5 commit 6b9e7b0

File tree

1 file changed

+9
-10
lines changed

1 file changed

+9
-10
lines changed

extension/module/module.h

Lines changed: 9 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,15 @@ class Module {
3838
MmapUseMlockIgnoreErrors,
3939
};
4040

41+
struct MethodHolder {
42+
std::vector<std::vector<uint8_t>> planned_buffers;
43+
std::vector<runtime::Span<uint8_t>> planned_spans;
44+
std::unique_ptr<runtime::HierarchicalAllocator> planned_memory;
45+
std::unique_ptr<runtime::MemoryManager> memory_manager;
46+
std::unique_ptr<runtime::Method> method;
47+
std::vector<runtime::EValue> inputs;
48+
};
49+
4150
/**
4251
* Constructs an instance by loading a program from a file with specified
4352
* memory locking behavior.
@@ -423,16 +432,6 @@ class Module {
423432
return event_tracer_.get();
424433
}
425434

426-
private:
427-
struct MethodHolder {
428-
std::vector<std::vector<uint8_t>> planned_buffers;
429-
std::vector<runtime::Span<uint8_t>> planned_spans;
430-
std::unique_ptr<runtime::HierarchicalAllocator> planned_memory;
431-
std::unique_ptr<runtime::MemoryManager> memory_manager;
432-
std::unique_ptr<runtime::Method> method;
433-
std::vector<runtime::EValue> inputs;
434-
};
435-
436435
private:
437436
std::string file_path_;
438437
LoadMode load_mode_{LoadMode::MmapUseMlock};

0 commit comments

Comments
 (0)