mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
6f2307ba6a
Summary: When using `setuptools` to build a Python extension, setuptools will automatically add an ABI suffix like `cpython-37m-x86_64-linux-gnu` to the shared library name when using Python 3. This is required for extensions meant to be imported as Python modules. When we use setuptools to build shared libraries not meant as Python modules, for example libraries that define and register TorchScript custom ops, having your library called `my_ops.cpython-37m-x86_64-linux-gnu.so` is a bit annoying compared to just `my_ops.so`, especially since you have to reference the library name when loading it with `torch.ops.load_library` in Python. This PR fixes this by adding a `with_options` class method to the `torch.utils.cpp_extension.BuildExtension` which allows configuring the `BuildExtension`. In this case, the first option we add is `no_python_abi_suffix`, which we then use in `get_ext_filename` (override from `setuptools.build_ext`) to throw away the ABI suffix. I've added a test `setup.py` in a `no_python_abi_suffix_test` folder. Fixes https://github.com/pytorch/pytorch/issues/14188 t-vi fmassa soumith Pull Request resolved: https://github.com/pytorch/pytorch/pull/14130 Differential Revision: D13216575 Pulled By: goldsborough fbshipit-source-id: 67dc345c1278a1a4ee4ca907d848bc1fb4956cfa
45 lines
1.1 KiB
Python
45 lines
1.1 KiB
Python
import argparse
|
|
import os.path
|
|
import sys
|
|
|
|
import torch
|
|
|
|
|
|
def get_custom_op_library_path():
|
|
if sys.platform.startswith("win32"):
|
|
library_filename = "custom_ops.dll"
|
|
elif sys.platform.startswith("darwin"):
|
|
library_filename = "libcustom_ops.dylib"
|
|
else:
|
|
library_filename = "libcustom_ops.so"
|
|
path = os.path.abspath("build/{}".format(library_filename))
|
|
assert os.path.exists(path), path
|
|
return path
|
|
|
|
|
|
class Model(torch.jit.ScriptModule):
|
|
def __init__(self):
|
|
super(Model, self).__init__()
|
|
self.p = torch.nn.Parameter(torch.eye(5))
|
|
|
|
@torch.jit.script_method
|
|
def forward(self, input):
|
|
return torch.ops.custom.op_with_defaults(input)[0] + 1
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(
|
|
description="Serialize a script module with custom ops"
|
|
)
|
|
parser.add_argument("--export-script-module-to", required=True)
|
|
options = parser.parse_args()
|
|
|
|
torch.ops.load_library(get_custom_op_library_path())
|
|
|
|
model = Model()
|
|
model.save(options.export_script_module_to)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|