mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-21 21:49:24 +08:00
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/30144 Create script to produce libtorch that only contains ops needed by specific models. Developers can use this workflow to further optimize mobile build size. Need keep a dummy stub for unused (stripped) ops because some JIT side logic requires certain function schemas to be existed in the JIT op registry. Test Steps: 1. Build "dump_operator_names" binary and use it to dump root ops needed by a specific model: ``` build/bin/dump_operator_names --model=mobilenetv2.pk --output=mobilenetv2.yaml ``` 2. The MobileNetV2 model should use the following ops: ``` - aten::t - aten::dropout - aten::mean.dim - aten::add.Tensor - prim::ListConstruct - aten::addmm - aten::_convolution - aten::batch_norm - aten::hardtanh_ - aten::mm ``` NOTE that for some reason it outputs "aten::addmm" but actually uses "aten::mm". You need fix it manually for now. 3. Run custom build script locally (use Android as an example): ``` SELECTED_OP_LIST=mobilenetv2.yaml scripts/build_pytorch_android.sh armeabi-v7a ``` 4. Checkout demo app that uses locally built library instead of downloading from jcenter repo: ``` git clone --single-branch --branch custom_build git@github.com:ljk53/android-demo-app.git ``` 5. Copy locally built libraries to demo app folder: ``` find ${HOME}/src/pytorch/android -name '*.aar' -exec cp {} ${HOME}/src/android-demo-app/HelloWorldApp/app/libs/ \; ``` 6. Build demo app with locally built libtorch: ``` cd ${HOME}/src/android-demo-app/HelloWorldApp ./gradlew clean && ./gradlew assembleDebug ``` 7. Install and run the demo app. In-APK arm-v7 libpytorch_jni.so build size reduced from 5.5M to 2.9M. Test Plan: Imported from OSS Differential Revision: D18612127 Pulled By: ljk53 fbshipit-source-id: fa8d5e1d3259143c7346abd1c862773be8c7e29a
94 lines
3.0 KiB
Python
94 lines
3.0 KiB
Python
import argparse
|
|
import os
|
|
import sys
|
|
|
|
source_files = {'.py', '.cpp', '.h'}
|
|
|
|
DECLARATIONS_PATH = 'torch/share/ATen/Declarations.yaml'
|
|
|
|
|
|
# TODO: This is a little inaccurate, because it will also pick
|
|
# up setup_helper scripts which don't affect code generation
|
|
def all_generator_source():
|
|
r = []
|
|
for directory, _, filenames in os.walk('tools'):
|
|
for f in filenames:
|
|
if os.path.splitext(f)[1] in source_files:
|
|
full = os.path.join(directory, f)
|
|
r.append(full)
|
|
return sorted(r)
|
|
|
|
|
|
def generate_code(ninja_global=None,
|
|
declarations_path=None,
|
|
nn_path=None,
|
|
install_dir=None,
|
|
subset=None,
|
|
disable_autograd=False,
|
|
selected_op_list_path=None):
|
|
# cwrap depends on pyyaml, so we can't import it earlier
|
|
root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
sys.path.insert(0, root)
|
|
from tools.autograd.gen_autograd import gen_autograd, gen_autograd_python
|
|
from tools.jit.gen_jit_dispatch import gen_jit_dispatch
|
|
|
|
# Build ATen based Variable classes
|
|
autograd_gen_dir = install_dir or 'torch/csrc/autograd/generated'
|
|
jit_gen_dir = install_dir or 'torch/csrc/jit/generated'
|
|
for d in (autograd_gen_dir, jit_gen_dir):
|
|
if not os.path.exists(d):
|
|
os.makedirs(d)
|
|
|
|
if subset == "pybindings" or not subset:
|
|
gen_autograd_python(declarations_path or DECLARATIONS_PATH, autograd_gen_dir, 'tools/autograd')
|
|
|
|
if subset == "libtorch" or not subset:
|
|
gen_autograd(
|
|
declarations_path or DECLARATIONS_PATH,
|
|
autograd_gen_dir,
|
|
'tools/autograd',
|
|
disable_autograd=disable_autograd,
|
|
)
|
|
gen_jit_dispatch(
|
|
declarations_path or DECLARATIONS_PATH,
|
|
jit_gen_dir,
|
|
'tools/jit/templates',
|
|
disable_autograd=disable_autograd,
|
|
selected_op_list_path=selected_op_list_path)
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description='Autogenerate code')
|
|
parser.add_argument('--declarations-path')
|
|
parser.add_argument('--nn-path')
|
|
parser.add_argument('--ninja-global')
|
|
parser.add_argument('--install_dir')
|
|
parser.add_argument(
|
|
'--subset',
|
|
help='Subset of source files to generate. Can be "libtorch" or "pybindings". Generates both when omitted.'
|
|
)
|
|
parser.add_argument(
|
|
'--disable-autograd',
|
|
default=False,
|
|
action='store_true',
|
|
help='It can skip generating autograd related code when the flag is set',
|
|
)
|
|
parser.add_argument(
|
|
'--selected-op-list-path',
|
|
help='Path to the yaml file that contains the list of operators to include for custom build.',
|
|
)
|
|
options = parser.parse_args()
|
|
generate_code(
|
|
options.ninja_global,
|
|
options.declarations_path,
|
|
options.nn_path,
|
|
options.install_dir,
|
|
options.subset,
|
|
options.disable_autograd,
|
|
options.selected_op_list_path,
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|