OSSForks/pytorch
1
0
Fork 0
mirror of https://github.com/zebrajr/pytorch.git synced 2025-12-07 00:21:07 +01:00
Code Issues Actions 38 Packages Projects Releases Wiki Activity
cuda_stream_jit
pytorch/torch/fx/experimental
History
Wang Xu 62d37b9f26 add size_based_partition final (#46282) 
Summary:
Reopen the PR: https://github.com/pytorch/pytorch/pull/45837
This PR add a new feature for Partitioner() class called size_based_partition. Given a list of devices with the same memory size, this function could distribute graph nodes into different devices. To implement this feature, several help functions are created in Partitioner.py and GraphManipulation.py.
An unit test is also added in test/test_fx_experimental.py

Pull Request resolved: https://github.com/pytorch/pytorch/pull/46282

Reviewed By: gcatron

Differential Revision: D24288470

Pulled By: scottxu0730

fbshipit-source-id: e81b1e0c56e34f61e497d868882126216eba7538
2020-10-14 03:44:05 -07:00
..
__init__.py
GraphManipulation.py add size_based_partition final (#46282) 2020-10-14 03:44:05 -07:00
Partitioner.py add size_based_partition final (#46282) 2020-10-14 03:44:05 -07:00
shape_prop.py [FX][WIP] Mutable Graph APIs (#45227) 2020-10-05 17:07:08 -07:00
subgraph_creation_example.py [FX] Make output a non-special Node (#45599) 2020-10-02 17:08:17 -07:00