Skip to content
Permalink
Browse files

Basic PyTorch Integration (#3069)

Summary:
Setup basics of model loader and import `mul` and `div` JIT IR nodes.
Still need to clean things up and figure out some things about the setup.py script.

Documentation:
Pull Request resolved: #3069

Test Plan:
Python source:
<img width="118" alt="Screen Shot 2019-06-10 at 8 26 09 AM" src="https://user-images.githubusercontent.com/1740091/59206260-6b4bf100-8b59-11e9-858d-583e317b168b.png">

JIT IR:
<img width="397" alt="Screen Shot 2019-06-10 at 8 25 14 AM" src="https://user-images.githubusercontent.com/1740091/59206203-4f484f80-8b59-11e9-94f9-1073f5ff278c.png">

Glow graph:
<img width="269" alt="Screen Shot 2019-06-10 at 8 24 46 AM" src="https://user-images.githubusercontent.com/1740091/59206204-4f484f80-8b59-11e9-96ab-ff591e50ec5b.png">

Results of running on JIT vs Glow match
<img width="347" alt="Screen Shot 2019-06-10 at 8 29 28 AM" src="https://user-images.githubusercontent.com/1740091/59206511-ef05dd80-8b59-11e9-9df7-4e99f05713ab.png">

Differential Revision: D15911449

Pulled By: jackm321

fbshipit-source-id: 3dbd2bb3819c254b987a93a64b39ca558f969d77
  • Loading branch information...
jackm321 authored and facebook-github-bot committed Jun 21, 2019
1 parent c115140 commit 5478d8719cca842374d2cd74f23bb4c554134aad
@@ -6,6 +6,10 @@
.DS_Store
*.so
*.dylib
*.egg-info
*.egg
*.eggs
*.pyc

GPATH
GRTAGS
@@ -16,3 +16,6 @@
[submodule "googlebenchmark"]
path = tests/googlebenchmark
url = https://github.com/google/benchmark
[submodule "thirdparty/pybind11"]
path = thirdparty/pybind11
url = https://github.com/pybind/pybind11.git
@@ -10,6 +10,7 @@ option(GLOW_WITH_LLVMIRCODEGEN "Build the LLVM-based code generation library" ON
option(GLOW_WITH_OPENCL "Build the OpenCL backend" OFF)
option(GLOW_WITH_HABANA "Build the Habana backend" OFF)
option(GLOW_BUILD_EXAMPLES "Build the examples" ON)
option(GLOW_BUILD_PYTORCH_INTEGRATION "Build integration for PyTorch" OFF)
option(GLOW_BUILD_TESTS "Build the tests" ON)
option(GLOW_WITH_BUNDLES "Build bundles" OFF)
option(LINK_PROTOBUF_AS_DLL "Link against protobuf build as dynamic libray." OFF)
@@ -176,6 +177,15 @@ if (GLOW_BUILD_EXAMPLES)
add_subdirectory(examples)
endif()

if (GLOW_BUILD_PYTORCH_INTEGRATION)
if(NOT EXISTS "${GLOW_THIRDPARTY_DIR}/pybind11")
message(FATAL_ERROR "No pybind11 git submodule. Run: git submodule update --init --recursive")
endif()

add_subdirectory(${GLOW_THIRDPARTY_DIR}/pybind11)
add_subdirectory(torch_glow/src)
endif()

if (GLOW_WITH_BUNDLES AND NOT GLOW_WITH_CPU)
message(FATAL_ERROR "Cannot create bundles without CPU backend. Configure with -DGLOW_WITH_BUNDLES and -DGLOW_WITH_CPU to build bundles.")
endif()
Submodule pybind11 added at a1b71d
@@ -0,0 +1,36 @@
import torch
import torch_glow

x = torch.randn(4)
y = torch.randn(4)

@torch.jit.script
def foo(a, b):
c = a.mul(b)
a = c.mul(c)
a = c.mul(a)
d = c.div(a)
return d

print("original jit ir")
print(foo.graph_for(x, y))

jit_res = foo(x, y)

torch_glow.enableFusionPass()

@torch.jit.script
def foo_glow(a, b):
return foo(a, b)

print("glow jit ir")
print(foo_glow.graph_for(x, y))

jit_glow_res = foo_glow(x, y)

print("jit_res")
print(jit_res)
print("jit_glow_res")
print(jit_glow_res)

assert torch.allclose(jit_res, jit_glow_res)
@@ -0,0 +1,5 @@
[aliases]
test=pytest

[tool:pytest]
testpaths = tests
@@ -0,0 +1,206 @@
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

from distutils.spawn import find_executable
from distutils import sysconfig, log
import setuptools
import setuptools.command.build_py
import setuptools.command.develop
import setuptools.command.build_ext

from collections import namedtuple
from contextlib import contextmanager
import glob
import os
import shlex
import subprocess
import sys
import argparse
from textwrap import dedent
import multiprocessing

try:
import torch
except ImportError as e:
print('Unable to import torch. Error:')
print('\t', e)
print('You need to install pytorch first.')
sys.exit(1)


FILE_DIR = os.path.realpath(os.path.dirname(__file__))
TOP_DIR = os.path.realpath(os.path.dirname(FILE_DIR))
CMAKE_BUILD_DIR = os.path.join(TOP_DIR, 'build')

CMAKE = find_executable('cmake') or find_executable('cmake3')
if not CMAKE:
print('Could not find "cmake". Make sure it is in your PATH.')
sys.exit(1)

install_requires = []
setup_requires = []
tests_require = []
extras_require = {}

# ################################################################################
# # Flags
# ################################################################################


parser = argparse.ArgumentParser()
parser.add_argument("--debug", type=bool, default=False, help="Compile with debug on")
parser.add_argument("--cmake_prefix_path", type=str, help="Populates -DCMAKE_PREFIX_PATH")
args = parser.parse_known_args()[0]


# ################################################################################
# # Utilities
# ################################################################################


@contextmanager
def cd(path):
if not os.path.isabs(path):
raise RuntimeError('Can only cd to absolute path, got: {}'.format(path))
orig_path = os.getcwd()
os.chdir(path)
try:
yield
finally:
os.chdir(orig_path)


# ################################################################################
# # Customized commands
# ################################################################################


class cmake_build(setuptools.Command):
"""
Compiles everything when `python setup.py develop` is run using cmake.
Custom args can be passed to cmake by specifying the `CMAKE_ARGS`
environment variable.
"""
def initialize_options(self):
pass

def finalize_options(self):
pass

def _run_cmake(self):
with cd(CMAKE_BUILD_DIR):
cmake_args = [
CMAKE,
'-DGLOW_BUILD_PYTORCH_INTEGRATION=ON',
'-DBUILD_SHARED_LIBS=OFF',
'-DCMAKE_EXPORT_COMPILE_COMMANDS=ON',
'-DCMAKE_BUILD_TYPE={}'.format('Debug' if args.debug else 'Release'),
'-DPYTHON_EXECUTABLE={}'.format(sys.executable),
# PyTorch cmake args
'-DPYTORCH_DIR={}'.format(
os.path.dirname(os.path.realpath(torch.__file__))),
]

if args.cmake_prefix_path:
cmake_args.append('-DCMAKE_PREFIX_PATH={}'.format(args.cmake_prefix_path))

if 'CMAKE_ARGS' in os.environ:
extra_cmake_args = shlex.split(os.environ['CMAKE_ARGS'])
log.info('Extra cmake args: {}'.format(extra_cmake_args))
cmake_args.extend(extra_cmake_args)
cmake_args.append(TOP_DIR)
subprocess.check_call(cmake_args)

def _run_build(self):
with cd(CMAKE_BUILD_DIR):
build_args = [
CMAKE,
'--build', os.curdir,
'--', '-j', str(multiprocessing.cpu_count()),
]
subprocess.check_call(build_args)

def run(self):
is_initial_build = not os.path.exists(CMAKE_BUILD_DIR)
if is_initial_build:
os.makedirs(CMAKE_BUILD_DIR)
self._run_cmake()
self._run_build()


class develop(setuptools.command.develop.develop):
def run(self):
self.run_command('build_ext')
setuptools.command.develop.develop.run(self)


class build_ext(setuptools.command.build_ext.build_ext):
def run(self):
self.run_command('cmake_build')
setuptools.command.build_ext.build_ext.run(self)

def build_extensions(self):
for ext in self.extensions:
fullname = self.get_ext_fullname(ext.name)
filename = os.path.basename(self.get_ext_filename(fullname))

src = os.path.join(CMAKE_BUILD_DIR, "torch_glow", "src", filename)
dst = os.path.join(os.path.realpath(self.build_lib), 'torch_glow', filename)
print("dst", dst)
if not os.path.exists(os.path.dirname(dst)):
os.makedirs(os.path.dirname(dst))
self.copy_file(src, dst)


cmdclass = {
'cmake_build': cmake_build,
'develop': develop,
'build_ext': build_ext,
}

# ################################################################################
# # Extensions
# ################################################################################

ext_modules = [
setuptools.Extension(
name=str('torch_glow._torch_glow'),
sources=[])
]

# ################################################################################
# # Packages
# ################################################################################

# # no need to do fancy stuff so far
packages = setuptools.find_packages()

# ################################################################################
# # Test
# ################################################################################

setup_requires.append('pytest-runner')
tests_require.append('pytest')

# ################################################################################
# # Final
# ################################################################################

setuptools.setup(
name="torch_glow",
description="PyTorch + Glow",
ext_modules=ext_modules,
cmdclass=cmdclass,
packages=packages,
include_package_data=True,
install_requires=install_requires,
setup_requires=setup_requires,
tests_require=tests_require,
extras_require=extras_require,
author='jackm321',
author_email='jackmontgomery@fb.com',
url='https://github.com/pytorch/glow',
)
@@ -0,0 +1,39 @@
# PYTORCH_DIR
if(DEFINED ENV{PYTORCH_DIR})
SET(PYTORCH_DIR $ENV{PYTORCH_DIR})
message(STATUS "Using PYTORCH_DIR from env")
endif()

if(NOT EXISTS "${PYTORCH_DIR}")
message(FATAL_ERROR "No PyTorch installation found")
endif()

message(STATUS "Using pytorch dir ${PYTORCH_DIR}")

link_directories(${PYTORCH_DIR}/lib)

pybind11_add_module(_torch_glow
binding.cpp
PyTorchModelLoader.cpp
CachingGraphRunner.cpp)

target_compile_options(_torch_glow PRIVATE -frtti -fexceptions)
target_link_libraries(_torch_glow
PRIVATE
# pytorch
torch
caffe2
c10
# glow
Support
ExecutionEngine
Graph
Importer
Interpreter
Backends
# pybind
pybind11)

target_include_directories(_torch_glow PUBLIC
${PYTORCH_DIR}/include
)

0 comments on commit 5478d87

Please sign in to comment.
You can’t perform that action at this time.