mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
cf3247b74a
Given an `package: _ExportPackage`, users can get a ready-to-use workspace in `tmp_dir` by calling:
```python
package._compiled_and_package(
tmp_dir + "/pt2_pacakge_name.pt2", True, package_example_inputs = True
)
```
`tmp_dir` will contains:
- `main.cpp` (an example cpp file that create the models, if package_example_inputs is True, it'll also load the example inputs and run the models)
- `CMakeLists.txt`
- `pt2_pacakge_name/` (this is where the models are)
- `pt2_pacakge_name.pt2`
- `inputs.pt` files if package_example_inputs is True
Remaining TODOs
- support loading contants/weights
- the `package_example_inputs = True` option only supports a list of Tensors for now
- eventually we should remove the `torch` dependency, and use `SlimTensor`/`StableIValue` instead.
Test Plan:
```
python test/inductor/test_aot_inductor_package.py -k test_compile_with_exporter
```
Example generated `main.cpp`:
```cpp
#include <dlfcn.h>
#include <fstream>
#include <iostream>
#include <memory>
#include <torch/torch.h>
#include <vector>
#include <torch/csrc/inductor/aoti_torch/tensor_converter.h>
#include "package/data/aotinductor/Plus__default/Plus__default.h"
#include "package/data/aotinductor/Minus__default/Minus__default.h"
using torch::aot_inductor::AOTInductorModelPlus__default;
using torch::aot_inductor::AOTInductorModelMinus__default;
using torch::aot_inductor::ConstantHandle;
using torch::aot_inductor::ConstantMap;
int main(int argc, char* argv[]) {
std::string device_str = "cpu";
try {
c10::Device device(device_str);
// Load input tensors for model Plus__default
std::vector<at::Tensor> input_tensors1;
for (int j = 0; j < 2; ++j) {
std::string filename = "Plus__default_input_" + std::to_string(j) + ".pt";
std::ifstream in(filename, std::ios::binary);
if (!in.is_open()) {
std::cerr << "Failed to open file: " << filename << std::endl;
return 1;
}
std::vector<char> buffer((std::istreambuf_iterator<char>(in)), std::istreambuf_iterator<char>());
torch::IValue ivalue = torch::pickle_load(buffer);
input_tensors1.push_back(ivalue.toTensor().to(device));
}
// Load input tensors for model Minus__default
std::vector<at::Tensor> input_tensors2;
for (int j = 0; j < 2; ++j) {
std::string filename = "Minus__default_input_" + std::to_string(j) + ".pt";
std::ifstream in(filename, std::ios::binary);
if (!in.is_open()) {
std::cerr << "Failed to open file: " << filename << std::endl;
return 1;
}
std::vector<char> buffer((std::istreambuf_iterator<char>(in)), std::istreambuf_iterator<char>());
torch::IValue ivalue = torch::pickle_load(buffer);
input_tensors2.push_back(ivalue.toTensor().to(device));
}
// Create array of input handles
auto input_handles1 =
torch::aot_inductor::unsafe_alloc_new_handles_from_tensors(input_tensors1);
auto input_handles2 =
torch::aot_inductor::unsafe_alloc_new_handles_from_tensors(input_tensors2);
// Create array for output handles
AtenTensorHandle output_handle1;
AtenTensorHandle output_handle2;
// Create and load models
auto constants_map1 = std::make_shared<ConstantMap>();
auto constants_array1 = std::make_shared<std::vector<ConstantHandle>>();
auto model1 = AOTInductorModelPlus__default::Create(
constants_map1, constants_array1, device_str,
"package/data/aotinductor/Plus__default/");
model1->load_constants();
auto constants_map2 = std::make_shared<ConstantMap>();
auto constants_array2 = std::make_shared<std::vector<ConstantHandle>>();
auto model2 = AOTInductorModelMinus__default::Create(
constants_map2, constants_array2, device_str,
"package/data/aotinductor/Minus__default/");
model2->load_constants();
// Run the models
torch::aot_inductor::DeviceStreamType stream1 = nullptr;
model1->run(&input_handles1[0], &output_handle1, stream1, nullptr);
torch::aot_inductor::DeviceStreamType stream2 = nullptr;
model2->run(&input_handles2[0], &output_handle2, stream2, nullptr);
// Convert output handles to tensors
auto output_tensor1 =
torch::aot_inductor::alloc_tensors_by_stealing_from_handles(&output_handle1, 1);
auto output_tensor2 =
torch::aot_inductor::alloc_tensors_by_stealing_from_handles(&output_handle2, 1);
// Validate outputs
std::cout << "output_tensor1" << output_tensor1 << std::endl;
std::cout << "output_tensor2" << output_tensor2 << std::endl;
return 0;
} catch (const std::exception &e) {
std::cerr << "Error: " << e.what() << std::endl;
return 1;
}
}
```
Rollback Plan:
Differential Revision: D78124705
Pull Request resolved: https://github.com/pytorch/pytorch/pull/158139
Approved by: https://github.com/desertfire
|
||
|---|---|---|
| .. | ||
| experimental | ||
| passes | ||
| pt2_archive | ||
| __init__.py | ||
| _draft_export.py | ||
| _remove_auto_functionalized_pass.py | ||
| _remove_effect_tokens_pass.py | ||
| _safeguard.py | ||
| _swap.py | ||
| _trace.py | ||
| _tree_utils.py | ||
| _unlift.py | ||
| _wrapper_utils.py | ||
| custom_obj.py | ||
| custom_ops.py | ||
| decomp_utils.py | ||
| dynamic_shapes.py | ||
| exported_program.py | ||
| graph_signature.py | ||
| unflatten.py | ||