pytorch

1
import argparse
2
import os.path
3
import sys
4
import torch
5

6

7
def get_custom_backend_library_path():
8
    """
9
    Get the path to the library containing the custom backend.
10

11
    Return:
12
        The path to the custom backend object, customized by platform.
13
    """
14
    if sys.platform.startswith("win32"):
15
        library_filename = "custom_backend.dll"
16
    elif sys.platform.startswith("darwin"):
17
        library_filename = "libcustom_backend.dylib"
18
    else:
19
        library_filename = "libcustom_backend.so"
20
    path = os.path.abspath(f"build/{library_filename}")
21
    assert os.path.exists(path), path
22
    return path
23

24

25
def to_custom_backend(module):
26
    """
27
    This is a helper that wraps torch._C._jit_to_test_backend and compiles
28
    only the forward method with an empty compile spec.
29

30
    Args:
31
        module: input ScriptModule.
32

33
    Returns:
34
        The module, lowered so that it can run on TestBackend.
35
    """
36
    lowered_module = torch._C._jit_to_backend("custom_backend", module, {"forward": {"": ""}})
37
    return lowered_module
38

39

40
class Model(torch.nn.Module):
41
    """
42
    Simple model used for testing that to_backend API supports saving, loading,
43
    and executing in C++.
44
    """
45

46
    def forward(self, a, b):
47
        return (a + b, a - b)
48

49

50
def main():
51
    parser = argparse.ArgumentParser(
52
        description="Lower a Module to a custom backend"
53
    )
54
    parser.add_argument("--export-module-to", required=True)
55
    options = parser.parse_args()
56

57
    # Load the library containing the custom backend.
58
    library_path = get_custom_backend_library_path()
59
    torch.ops.load_library(library_path)
60
    assert library_path in torch.ops.loaded_libraries
61

62
    # Lower an instance of Model to the custom backend  and export it
63
    # to the specified location.
64
    lowered_module = to_custom_backend(torch.jit.script(Model()))
65
    torch.jit.save(lowered_module, options.export_module_to)
66

67

68
if __name__ == "__main__":
69
    main()
70