Terra: Imperative-Symbolic Co-Execution of Imperative Deep Learning Programs

Imperative programming allows users to implement their deep neural networks ( DNNs ) easily and has become an essential part of recent deep learning ( DL ) frameworks . Recently , several systems have been proposed to combine the usability of imperative programming with the optimized performance of symbolic graph execution . Such systems convert imperative Python DL programs to optimized symbolic graphs and execute them . However , they can not fully support the usability of imperative programming . For example , if an imperative DL program contains a Python feature with no corresponding symbolic representation ( e.g. , third-party library calls or unsupported dynamic control flows ) they fail to execute the program . To overcome this limitation , we propose Terra , an imperative-symbolic co-execution system that can handle any imperative DL programs while achieving the optimized performance of symbolic graph execution . To achieve this , Terra builds a symbolic graph by decoupling DL operations from Python features . Then , Terra conducts the imperative execution to support all Python features , while delegating the decoupled operations to the symbolic execution . We evaluated Terra ’ s performance improvement and coverage with ten imperative DL programs for several DNN architectures . The results show that Terra can speed up the execution of all ten imperative DL programs , whereas AutoGraph , one of the state-of-the-art systems , fails to execute five of them . ∗Work done at Seoul National University †Corresponding author 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . 1 Introduction . The rapid evolution of deep neural networks ( DNNs ) has been fueled by the support of deep learning ( DL ) frameworks [ 6 , 32 ] . Such frameworks provide users with a programming layer to build and execute DNNs , commonly adopting Python as their host language . Typically , they execute DL programs with one of the two execution models : imperative or symbolic execution . In the former , the Python interpreter executes a DL program as a normal program , invoking DL operations on-the-fly . The invoked DL operations are executed on a separate DL accelerator asynchronously , and the Python interpreter continues running the program . The dynamic control flows of the DL operations are naturally expressed by the interpretation of the program , and users can utilize any functionalities of Python ( e.g. , dynamic typing and third-party libraries [ 9 , 16 ] ) while executing DL operations . On the other hand , in the latter model , the Python interpreter embeds DL operations into a symbolic graph that represents the entire dataflow of a DNN . Thus , users should define their DL programs only with existing symbolic operations that DL frameworks support . In other words , the dynamic control flows of a DNN should be explicitly represented by control flow operations ( e.g. , tf.cond and tf.while of TensorFlow ) . However , the symbolic execution can take advantages of various optimization techniques because the symbolic graph contains a whole computation lineage of a DNN architecture . Although symbolic execution achieves higher performance compared to imperative execution , imperative execution has been preferred because of its usability . Several systems [ 2 , 3 , 4 , 15 , 19 , 29 , 39 ] have been proposed to match the speed of symbolic execution while enjoying the benefit of imperative execution . These systems attempt to generate a symbolic graph that represents an entire imperative program and execute the graph instead of imperatively running the program . Methods for generating the symbolic graph can be broadly classified into two approaches : single path tracing and static compilation . The former approach generates a symbolic graph by imperatively executing a single iteration of a program and recording the executed DL operations . Systems that adopt the latter approach translate the abstract syntax tree ( AST ) of a program into a symbolic graph . Unfortunately , both approaches can correctly handle only a subset of imperative DL programs . For example , dynamic control flows in an imperative program are not captured by the single path tracing approach . On the other hand , the static compilation approach can not correctly generate a symbolic graph if a target program contains an AST node that does not have a corresponding symbolic operation such as try-excepts , generators , Python object mutations , and third-party library calls . As a result , it is up to the users to clearly understand the limited usability of these systems . In this paper , we propose Terra , an imperative-symbolic co-execution system that addresses the limitations . While the previous approaches replace the imperative execution with the symbolic execution , Terra maintains the imperative execution to support all Python features where DL operations are delegated to the symbolic execution . Also , Terra generates a symbolic graph by utilizing multiple traces of an imperative program . To be specific , Terra imperatively runs an imperative DL program for several iterations and collects traces , each of which is a linear chain of DL operations sorted by the execution order . The collected traces are merged into a TraceGraph , a directed acyclic graph ( DAG ) that encapsulates the captured DL operations along with their diverse execution orders . Terra stops collecting traces when the trace of the latest iteration is already embedded in the TraceGraph . To generate an executable symbolic graph from the TraceGraph , Terra adds additional information to the TraceGraph . First of all , Terra annotates the TraceGraph to enable communication between the imperative execution and the symbolic execution . In addition , Terra further analyzes the TraceGraph to insert control flow operations explicitly , so that a symbolic graph can be executed with the DL operations in a certain trace . After generating a symbolic graph from the TraceGraph , Terra starts the co-execution of a skeleton imperative program , in which DL operations are not performed , and the symbolic graph that represents the DL operations . Here , Terra continually checks whether the current trace is being expressed by the TraceGraph . If Terra detects a new trace , Terra discards the symbolic graph and collects more traces by running the original program imperatively . Terra then re-generates the symbolic graph and restarts the co-execution . Consequently , Terra is able to run any imperative DL programs correctly and efficiently even if it contains the Python features that the previous approaches can not handle . We have implemented Terra on TensorFlow v2.4.1 and compared Terra with TensorFlow ’ s imperative execution [ 7 ] and AutoGraph [ 29 ] . Our evaluation shows that Terra can train ten imperative DL programs including convolutional neural networks , transformer-based networks , and generative adversarial networks up to 1.73x faster than the original imperative execution . However , AutoGraph fails to support five programs for three reasons : third-party library call , Python object mutation , and tensor materialization during conversion , which we describe in § 2.2 . 2 Background & Related Works . 2.1 Imperative and Symbolic Execution . The imperative execution ( a.k.a . define-by-run ) [ 30 , 32 , 43 ] treats a DL program entirely as a typical Python program . Whenever the Python interpreter encounters a statement that declares a DL operation ( e.g. , z = torch.matmul ( x , y ) ) , the interpreter asynchronously invokes a corresponding computation kernel of a DL accelerator ( typically GPU or TPU ) . The imperative execution highly improves the programmability of DL programs because users can fully utilize the convenient language features and rich ecosystem of Python , including built-in functions , dynamic control flows , dynamic typing , and third-party libraries . However , since the imperative execution can not obtain a whole view of DNN computation , it misses optimization opportunities that the symbolic execution explained below can carry out . The symbolic execution ( a.k.a . define-and-run ) [ 6 , 14 , 42 ] executes a pre-built symbolic graph . In the symbolic execution , users should express their DNN architectures as symbolic graphs using the three kinds of symbolic operations : DL operations , control flow operations , and auxiliary operations . The DL operations are conventional compute-intensive operations ( e.g. , matrix multiplication or convolutional layer ) , and the control flow operations determine which DL operations to be executed based on a tensor value within the graph . Finally , to mitigate the limited usability of the symbolic execution , it supports auxiliary operations ( e.g. , tf.print and tf.py_function of TensorFlow ) . After the symbolic graph is constructed , graph optimizations could be applied such as operation fusion [ 11 , 20 , 26 , 31 , 35 , 41 , 46 , 47 , 48 ] , parallelized execution [ 23 , 31 , 35 ] , device placement [ 27 , 28 , 48 ] , layout optimization [ 20 , 25 ] , and memory optimization [ 8 , 17 , 18 ] . Then , an optimized graph executor such as TVM [ 11 ] , TensorRT [ 31 ] , TFRT [ 40 ] , and XLA [ 41 ] undertakes the actual execution of the symbolic graph . 2.2 Imperative Program with Symbolic Execution . There are two approaches to gain the usability of imperative execution and the performance of symbolic execution simultaneously . They attempt to convert an imperative DL program to a symbolic graph and exploit the symbolic execution with the converted graph . The single path tracing approach ( e.g. , torch.trace [ 3 ] , JAX [ 15 ] , and tf.function [ 39 ] ) executes an imperative DL program once and records all DL operations that were executed . A single linear chain of the executed DL operations , which is called a trace , becomes a symbolic graph of the imperative program . The symbolic graph is executed instead of the imperative program for subsequent iterations . Although the single path tracing approach looks very simple and intuitive , it is hard to capture dynamic control flows of the imperative program with this approach . To reflect them correctly , users need to explicitly declare the control flow operations for all dynamic control flows in their imperative programs , which undermines the programmability of the imperative program . Moreover , Python features that do not have corresponding symbolic operations ( e.g. , mutation of Python objects , use of third-party libraries ) are neither captured by the trace . It can yield an incorrect result since , in the following iterations , a graph executor executes a symbolic graph , which does not contain the Python features . The static compilation approach ( e.g. , TorchScript [ 4 ] and JANUS [ 19 ] ) is proposed to resolve the problem of the single path tracing approach . In contrast to the single path tracing approach , the static compilation approach does not extract a trace to generate a symbolic graph . Instead , it traverses an abstract syntax tree ( AST ) of the imperative program and directly converts each AST node to a corresponding symbolic operation . Even if it guarantees the correctness of the program , it fails to convert a program to a symbolic graph if the program contains a Python feature with no corresponding symbolic operation . We classify the representative failure cases into four categories : third-party library call , tensor materialization during conversion , dynamic control flow , and Python object mutation . Figure 1 presents simple examples of the cases . The static compilation approach can not convert the third-party library calls of Figure 1a and fails when it attempts to materialize the tensor data ( loss.numpy ( ) ) during the conversion . Figure 1b shows the generator in Python that the static compilation approach can not convert . For Figure 1c , only JANUS handles it by implementing custom auxiliary operations ( i.e. , GetAttr and SetAttr operations ) that access the Python heap during the symbolic execution . To increase the practicality of the previous approaches , AutoGraph [ 29 ] combines the static compilation approach with the single path tracing approach . AutoGraph converts AST nodes that correspond to dynamic control flow such as if-else , for , and while to new AST nodes representing proper control flow operations such as tf.cond and tf.while . Then , AutoGraph generates a symbolic graph by applying the single path tracing approach to the converted AST . Unfortunately , it can not fully support various kinds of dynamic control flows such as generator and try-except of Python . Hence , AutoGraph also entails the same correctness problem of the single path tracing approach when the imperative DL program employs the features described in Figure 1 . When the static compilation approach detects unsupported features , it just raises an error and aborts the program execution . To avoid this , users have to write their programs using only the supported features or put in additional efforts such as annotating types and refactoring functions to use tensor objects for their inputs and outputs . In this regard , AutoGraph [ 29 ] and TorchScript [ 4 ] provide the official language references [ 1 , 5 ] that users should be aware of before writing a target imperative program . However , those references usually require expert knowledge to understand , imposing a steep learning curve for new users . For example , the language document of AutoGraph [ 1 ] spans roughly twelve pages and divides features that AutoGraph does not support into four categories and ten subcategories . TorchScript provides an official language specification [ 5 ] by enumerating supported features over eleven pages . LazyTensor [ 36 ] is a concurrent work related to ours . It uses the lazy evaluation approach , in which the Python interpreter and symbolic graph executor run alternately . The Python interpreter executes an imperative DL program as it is and extracts a linear trace of operations . LazyTensor then checks whether the extracted trace is already cached or not . If cached , LazyTensor directly executes the cached graph . If not , it compiles and executes the new graph , then stores the graph for further executions . With the lazy evaluation , LazyTensor could support all Python features while executing the symbolic graph as Terra does . However , it has an unavoidable performance overhead due to the alternate execution . In other words , the graph executor progresses only when the Python interpreter finishes checking the existence of the cached graph . Similarly , the Python interpreter progresses after finishing the graph execution . Moreover , if an imperative program has a dynamic control flow that is not determined by a tensor value , LazyTensor fails to capture the control flow transparently . In this case , LazyTensor would work inefficiently because the traces could be different for each iteration .

Summary: This paper proposes a method for obtaining a symbolic (graph) representation of a model, given the imperative Python code. It distinguishes itself from (most) other methods by co-executing the symbolic as well as the imperative representation, as well as merging multiple traces into a single symbolic representation. This allows it to obtain arbitrary representations of