Computational graph optimization

Abstract

Methods, systems, and apparatus for optimizing the execution of operations of a neural network, including a computer program encoded on a computer storage medium. One method includes: obtaining data representing a graph characterizing a plurality of operations of a neural network, wherein each node of the graph characterizes an operation of the neural network and each edge of the graph characterizes a data dependency between the operations; processing the data representing the graph using a graph embedding neural network to generate an embedding of the graph; and processing the embedding of the graph using a policy neural network to generate a corresponding task output for each of a plurality of optimization tasks, wherein for each of the plurality of operations of the neural network, each task output includes a corresponding decision for the corresponding optimization task.

Description

Technical Field

[0001] This specification relates to the use of neural networks to process computational graphs. Background Technology

[0002] A neural network is a machine learning model that uses one or more non-linear units to predict the output of a received input. Some neural networks include one or more hidden layers in addition to an output layer. The output of each hidden layer is used as the input to one or more other layers in the network (i.e., one or more other hidden layers, an output layer, or both). Summary of the Invention

[0003] This specification describes a system implemented as a computer program on one or more computers at one or more locations, which receives data characterizing a neural network as input and generates an execution optimization strategy for optimizing the execution of operations of the neural network on one or more processing devices. In this specification, for each operation of the neural network, the "execution optimization strategy" for the neural network is data specifying a corresponding decision for each of one or more different execution optimization tasks.

[0004] The system can represent a neural network using a computational graph (also referred to as a "graph" for simplicity), which includes nodes, at least some of which are connected by edges. Each node in the graph represents an operation of the neural network, and each edge represents a data dependency between operations of the neural network. Using this graph, the system can generate an execution optimization strategy for the neural network. Then, one or more processing devices can execute the operations of the neural network using the decisions for each execution optimization task.

[0005] The subject matter described in this specification can be implemented in certain embodiments to achieve one or more of the following advantages.

[0006] The techniques described in this specification allow for the optimization of the execution of neural network operations by generating policies for one or more optimization tasks. Some of the policy generation systems described in this specification can be trained such that they are invariant to the underlying topology of the graph, thus enabling the generation of policies for a wide set of input graphs. Therefore, this architecture can be well generalized to previously unseen graphs.

[0007] Some of the techniques described in this specification allow for the generation of policies for multiple different execution optimization tasks at once. Because graph execution optimization problems can be strongly coupled, joint optimization for multiple different execution optimization tasks can lead to solutions that improve the runtime and device utilization of neural networks.

[0008] Some existing systems rely on manual tuning heuristics to optimize the performance of specific neural networks. These manual tuning heuristics can be device-specific and therefore cannot be generalized to other devices. Furthermore, the heuristics need to be continuously modified and fine-tuned using expert domain knowledge to accommodate new cases arising from previously unseen model architectures. The techniques described in this specification are flexible for a wide range of devices and model architectures, including those not encountered during the training of the neural network.

[0009] Some existing systems aim to learn execution-optimized policies for performing operations on a single, specific neural network. That is, the existing system is retrained for each individual neural network, and the training parameters of the existing system are not transferable to generate policies for other neural networks. Using the techniques described in this specification, it is possible to train a policy generation system to generate execution-optimized policies for a wide variety of different neural networks. These techniques can be applied to computational graphs with drastically different architectures. Furthermore, these techniques enable the generation of execution-optimized policies for graphs with realistic workloads (e.g., graphs with more than 1M or 10M nodes).

[0010] Some existing techniques for learning execution optimization strategies for performing operations on neural networks are computationally expensive, for example, generating decisions for a single node per iteration of the neural network. The system described in this specification can generate execution optimization strategies for a specific neural network in a single pass, thereby significantly reducing the time required to generate such execution optimization strategies.

[0011] Details of one or more embodiments of the subject matter of this specification are set forth in the accompanying drawings and the following description. Other features, aspects, and advantages of the subject matter will become apparent from the specification, drawings, and claims. Attached Figure Description

[0012] Figure 1 This is a diagram of an example policy generation system.

[0013] Figure 2 This is a graph of an example policy neural network.

[0014] Figure 3 This is a diagram of the example task subnet.

[0015] Figure 4 This is a flowchart used to generate an example process for implementing optimization strategies.

[0016] Figure 5 This is a flowchart of an example process for generating multi-task output.

[0017] The same reference numerals and names in the various figures indicate the same elements. Detailed Implementation

[0018] This specification describes a system for generating an execution optimization strategy for optimizing the operations of a neural network on one or more processing devices. The operations of the neural network can include those performed by the one or more processing devices during the training of the neural network. Alternatively or additionally, the operations of the neural network can include those performed by the one or more processing devices during the inference time after the neural network has been trained. The system is capable of generating task outputs for each of the one or more execution optimization tasks. For each operation of the neural network, the task output for the specific execution optimization task can include a corresponding decision for the specific execution optimization task.

[0019] The system can represent the operations of a neural network as a computation graph, where each node of the graph represents an operation of the neural network, and each edge of the graph represents the data dependency between the operations of the neural network.

[0020] For example, one of the tasks for performing the optimization task can be device placement, where each operation of the neural network is assigned to be performed on a specific one of a plurality of processing devices. The device can include any suitable type of computer hardware device, i.e., any device capable of performing at least some of the operations represented in the computation graph. In some implementations, the devices are heterogeneous. For example, the device can include a combination of any of the following: a central processing unit (CPU), a graphics processing unit (GPU), a tensor processing unit (TPU), other application-specific integrated circuits (ASICs) or other specialized hardware, a field-programmable gate array (FPGA), etc. In some other implementations, the devices are homogeneous, i.e., including only devices of the same type, i.e., only devices of one of the aforementioned types or only devices consisting of the same combination of the aforementioned types. Therefore, for each node of the computation graph, the task output for performing the optimization task is an identifier of the specific processing device to which the operation represented by that node is assigned, such as a one-hot encoding identifying the processing device.

[0021] As another example, one of the optimization tasks could be operation scheduling, where the system generates a schedule for performing operations on a neural network. As a specific example, the system could use a priority-based scheduler, where each operation assigned to a particular device is placed in a priority queue when it is ready to be executed (i.e., when all inputs to the operation are available). Here, the task output for this optimization task could be an identifier of the priority value for each operation of the neural network, such as a one-hot encoding that identifies the priority value. Then, when operations are pushed and popped from the priority queue, the processing device assigned to execute the operation can use the corresponding priority value for each operation of the neural network.

[0022] As another example, one of the optimization tasks can be operation fusion, where multiple operations of a neural network are combined such that a processing device assigned to perform multiple operations treats them as a single operation. For example, operation fusion can be useful when the output of a first operation is the input of a second operation. In this case, if the processing device treats the two operations as a single operation, the device does not need to write the output of the first operation to memory and then read the output back from memory later when the device is performing the second operation. Therefore, the task output for this optimization task can be the identifiers of two or more nodes in the computation graph that are to be treated as a single operation by the assigned processing device.

[0023] These execution optimization tasks are strongly coupled, meaning that decisions made for one execution optimization task can influence the optimal decision for another. As a specific example, if the system considers decisions made during operation fusion when placing execution devices, it can assign two operations that are good candidates for fusion to the same device. As another specific example, if the system considers decisions made during operation scheduling when placing execution devices, it can evenly distribute high-priority operations across all devices and low-priority operations across all devices, resulting in high device utilization. In other words, when multiple different execution optimization tasks are considered jointly, the system can leverage knowledge sharing across multiple tasks to generate optimal task outputs for each of the multiple tasks.

[0024] Figure 1 This is a diagram of an example policy generation system 100. Policy generation system 100 is an example of a system implemented as a computer program on one or more computers at one or more locations, wherein the systems, components, and techniques described below are capable of implementation.

[0025] The strategy generation system 100 is able to receive operation data 102 representing the operations of a neural network and process the operation data 102 to generate an execution optimization strategy 132, the execution optimization strategy 132 including a corresponding decision for each operation of the neural network for each of one or more execution optimization tasks.

[0026] The policy generation system 100 may include a graph representation subsystem 110, a graph embedding neural network 120, and a policy neural network 130.

[0027] The diagram shows that the subsystem 110 can receive operation data 102 as input. For example, for each operation of the neural network, the operation data 102 can include data representing: i) the type of operation, ii) a list of other computations of the neural network whose output is the input of the operation, iii) a list of other operations of the neural network that take the output of the operation as input, and iv) the shape of the output of the operation. The shape of the output of the operation can be the dimension of the output tensor of the operation, for example, specifying data that the output tensor has a dimension of 10 × 20 × 5 elements.

[0028] The graph representation subsystem 110 can process operation data 102 to generate graph data 112, which represents a graph characterizing the operations of a neural network. The graph includes multiple nodes, with at least some nodes connected by edges. Each node in the graph can characterize a different operation of the neural network, and each edge of the graph can characterize a data dependency between corresponding operations. That is, if the output of a first operation of the neural network is the input of a second operation of the neural network, then an edge can exist between a first node in the graph representing the first operation and a second node in the graph representing the second operation. In some embodiments, the graph is directed; that is, each edge of the graph has an associated direction such that one incident node of the edge is a source and the other incident node of the edge is a sink. In some other embodiments, the graph is undirected, where each edge does not have an associated direction.

[0029] Each node of the graph can be represented in the graph data 112 by initial node embeddings. In some implementations, the initial node embeddings for each node of the graph are machine learning embeddings of the features of the corresponding operations, i.e., generated by processing the features of the nodes using a machine learning model.

[0030] In some other implementations, the initial node embedding for each node of the graph is a combination of features corresponding to the operation, as described in operation data 102. For example, each initial node embedding can be a concatenation of the following: i) an embedding of the type of the corresponding operation, ii) an identifier of each adjacent node in the graph, and iii) an embedding of the shape of the output of the corresponding operation. The embedding of the type of the corresponding operation can be a one-hot encoding of the type of operation; that is, the embedding can have the same number of elements as the number of possible operation types, where the element corresponding to the type of the corresponding operation is "1" and all other elements are "0". The embedding of the shape of the output of the corresponding operation can have corresponding elements for each dimension of the output, for example, three elements corresponding to the length, height, and depth of the output, respectively. The value in each element of the shape embedding of the output can be the size of the output in the corresponding dimension; for example, the embedding (10,5,3) can indicate that the output has 10 elements in length, 5 elements in height, and 3 elements in depth. In some implementations, the identifier of each neighboring node of a specific node can be an embedding including elements corresponding to each node in the graph, wherein each element is "1" if the corresponding node is a neighbor of the specific node, and "0" if the corresponding node is not a neighbor of the specific node. In some other implementations, the identifier of each neighboring node of a specific node can be a list of integers, wherein each integer corresponds to a specific neighboring node in the graph. In some such implementations, the graph representation subsystem 110 imposes a fixed upper limit on the number of neighboring nodes of each node to ensure that each initial node embedding is of fixed length.

[0031] In this specification, an embedding is an ordered set of numerical values ​​representing inputs in a specific embedding space. For example, an embedding can be a vector of floating-point or other numerical values ​​with a fixed dimension.

[0032] Graph data can also include an adjacency matrix for the nodes of the graph. The adjacency matrix for a graph with n nodes is an n×n matrix, where the (i,j)th element is 1 if there is an edge between node i and node j, and 0 otherwise.

[0033] The graph embedding neural network 120 can take graph data 112 as input. In some embodiments, the policy generation system 100 can receive graph data 112 instead of operation data 102 as input; that is, in some embodiments, the policy generation system 100 does not include a graph representation subsystem 110.

[0034] The graph embedding neural network 120 can process graph data 112 to generate graph embeddings 122. Graph embedding 122 can have n elements, where each element is a node embedding used for updating the nodes of the graph. For example, graph embedding 122 can be a matrix with n rows, where each row is the updated node embedding of the corresponding node of the graph.

[0035] In some such implementations, the graph embedding neural network 120 is capable of generating updated embeddings for each node of the graph across multiple embedding time steps. That is, at each embedding time step, for each node of the graph, the graph embedding neural network 120 generates a new embedding for that node by processing the embeddings of the nodes of the graph generated in the previous embedding time step. In the first embedding time step, the graph embedding neural network 120 is capable of generating new node embeddings by processing the initial node embeddings in the graph data 112.

[0036] At a given embedding time step, the graph embedding neural network 120 is able to first generate neighborhood embeddings of nodes representing the neighbors of each node in the graph. If two nodes are connected by an edge, the first node is a neighbor of the second node. In some embodiments, the graph embedding neural network 120 is able to compute the average of the embeddings of neighboring nodes generated in previous embedding time steps. In some other embodiments, the graph embedding neural network 120 is able to use a combined neural network to process the embeddings of neighboring nodes generated in previous embedding time steps. As a specific example, the graph embedding neural network 120 is able to compute a neighborhood embedding h for each node v. N(v) :

[0037]

[0038] in, It is the embedding of the neighboring nodes of v generated in the previous embedding time step, where N(v) is the set of the neighboring nodes of v, (W (l) b (l) Define an affine transformation, where σ is the sigmoid activation function and max is the max pooling layer.

[0039] Then, for each node in the graph, the graph embedding neural network 120 can combine i) the node's neighborhood embedding and ii) the node's embedding generated at a previous time step to generate a new embedding for the node at the embedding time step. In some implementations, for each node in the graph, the graph embedding neural network 120 can cascade two embeddings to generate a combined embedding for the node. Then, for each node in the graph, the graph embedding neural network 120 can generate a new embedding for the node by processing the combined embedding of the node (e.g., by processing the combined embedding of nodes with one or more fully connected neural network layers).

[0040] Graph embedding neural network 120 can provide graph embedding 122 to policy neural network 130. Policy neural network 130 can process graph embedding 130 to generate execution optimization policy 132. Execution optimization policy 132 can include corresponding task outputs for each of one or more execution optimization tasks. For each operation of the neural network, the task output for a specific execution optimization task includes a decision for that specific execution optimization task.

[0041] In some implementations, the graph embedding neural network 120 and the policy neural network 130 can be trained end-to-end jointly to optimize a specific objective function. That is, the training system can backpropagate errors through each neural network, thereby determining parameter updates for each neural network in parallel. In some other implementations, one of the neural networks can be "frozen" while training the other neural network. That is, the system does not determine parameter updates for the "frozen" neural network during backpropagation.

[0042] In some implementations, the training system can use an objective function to train the policy generation system 100, which measures the performance of the neural network's operations when performing operations according to the generated execution optimization policy 132. As a specific example, the training system can train the policy generation system 100 to minimize a metric of peak memory usage or a metric of average memory usage.

[0043] As another specific example, the training system can train the policy generation system 100 using an objective function representing the expected runtime of operations on a set of N candidate computation graphs. Each candidate computation graph represents a neural network whose operations can be optimized using the policy generation system 100. For example, the N candidate computation graphs can include a training set of computation graphs used during the training of the policy generation system 100. Therefore, by training the policy generation system 100 to minimize the objective function, the training system can train the policy generation system 100 to reduce the expected runtime of operations on the N candidate graphs. In some such implementations, the system uses the following objective function J(θ):

[0044] Where G is a candidate graph. It is the space of candidate graphs, where T represents the task output generated by policy generation system 100 for one or more execution optimization tasks, and π θ (G) represents the strategy for generating a candidate graph G using the current parameter values ​​θ of the strategy generation system 100, and r G,TThe reward is a function of the runtime of the operation represented by the candidate graph G using the task output T. The training system is able to update the parameter values ​​of the policy neural network 130 with the objective function J(θ) after generating the task output T for one or more tasks to perform optimization, and then backpropagate the update to the graph embedding neural network 120.

[0045] In some implementations, the training system can use proximity policy optimization to optimize the objective function.

[0046] In some implementations, the training system can evaluate the generated execution optimization policy on a real system, for example, by executing the neural network on one or more real processing devices and measuring the runtime. Alternatively or additionally, the training system can use a model that estimates the runtime of the neural network's operations to process the task output generated by the policy generation system 100.

[0047] Figure 2 This is a diagram of an example policy neural network 200. The policy neural network 200 is an example of a system implemented as a computer program on one or more computers at one or more locations, wherein the systems, components, and techniques described below are capable of implementation.

[0048] The policy neural network 200 receives node embeddings 202 of the nodes of a graph representing the operations of the neural network as input, and generates corresponding task outputs 222a-c for each of one or more optimization tasks. The node embeddings 202 may have been derived from a graph embedding neural network (e.g., Figure 1 The graph depicted in the diagram is generated by an embedded neural network 120.

[0049] The policy neural network 200 can include a policy embedding subnetwork 210 and corresponding task subnetworks 220a-c for each of one or more optimization tasks. The policy embedding subnetwork 210 can receive node embeddings 202 and process the node embeddings 202 using one or more attention neural network layers (e.g., neural network layers that process the input using multi-head attention).

[0050] In some implementations, the attention neural network layers are bidirectional. That is, for each node embedding in node embedding 202, each attention neural network layer generates an output corresponding to the node embedding by processing the output generated by the previous attention neural network layer corresponding to each node embedding in node embedding 202.

[0051] In some other implementations, the attention neural network layers are unidirectional or "masked." That is, for each particular node embedding in node embedding 202 and given the order of node embeddings 202, each attention neural network layer generates an output corresponding to a node embedding by processing the output generated by the previous attention neural network layer that corresponds to the particular node embedding and the previous node embedding in the order of node embeddings 202. In some implementations where the graph is undirected, the node embeddings 202 can be ordered according to the topological ordering of the nodes of the graph. The topological ordering of a graph with directed edges and no cycles is a linear ordering of the nodes of the graph such that for each directed edge between a first node and a second node, the first node is earlier in the order than the second node. In some other implementations, the node embeddings 202 can be ordered randomly.

[0052] In some such implementations, the attention neural network layer uses a segmented recursive mechanism. That is, the node embedding 202 is segmented into multiple segments, where each node embedding is assigned to a specific segment within a specific segment sequence. For each segment and each attention neural network layer, the output of the attention neural network layer corresponding to each node embedding in the segment can be cached, and these cached outputs can be used as input to the next attention neural network layer as the next attention neural network layer is generating an output corresponding to a node embedding in the next segment of the segment sequence. In other words, for each specific node embedding, the input to the attention neural network layer can include i) the output corresponding to the node embedding generated by a previous attention neural network layer, ii) the output corresponding to a previous node embedding in the same segment generated by a previous attention neural network layer, and iii) one or more cached outputs corresponding to the corresponding node embedding in a previous segment of the segment sequence generated by a previous attention neural network layer.

[0053] In an embodiment where the policy neural network 200 generates a task output for a single execution optimization task, the output of the policy embedding subnetwork 210 can be the task output for the execution optimization task. That is, the policy neural network 200 includes only the policy embedding subnetwork 210. For example, the policy embedding subnetwork 210 can include the aforementioned attention neural network layer and an output neural network layer, which processes the output of the last attention neural network layer corresponding to the node embedding of a specific node in the graph, and generates a decision for the execution optimization task for the operation represented by the specific node.

[0054] In an implementation where the policy neural network 200 generates task outputs for each of a plurality of execution optimization tasks, the output of the policy embedding subnetwork 210 can be a first task embedding 212a. Then, the policy neural network 200 can sequentially generate task outputs for each of the plurality of execution optimization tasks. When generating task outputs for a corresponding execution optimization task, the policy neural network 200 can use information about the task outputs for other execution optimization tasks that have already been generated.

[0055] Specifically, the first task subnet 220a is capable of processing the first task embedding 212a to generate a first task output 222a corresponding to a first execution optimization task among a plurality of execution optimization tasks. The first task subnet 220a is also capable of generating a second task embedding 212b that includes information about the first task output 222a and the first task embedding 212a.

[0056] The second task subnet 220b is capable of processing the second task embedding 212b to generate a second task output 222b corresponding to a second execution optimization task among a plurality of execution optimization tasks. The second task subnet 220b is also capable of generating a third task embedding 212c that includes information about the second task output 222b and the second task embedding 212b.

[0057] The sequence continues until the final task subnet 220m, corresponding to the m-th and last optimized task, which processes the task embeddings generated by the previous task subnets to generate the m-th task output 222m. The final task subnet 220m does not generate task embeddings because there are no further tasks. Therefore, the i-th task subnet uses information from all previous task outputs to generate the i-th task output by processing the i-th task embedding, which is generated from the (i-1)-th task output and the (i-1)-th task embedding. The (i-1)-th task embedding itself is generated from the (i-2)-th task output and the (i-2)-th task embedding, and so on.

[0058] In some implementations, the policy neural network 200 can be tuned based on features of the computation graph represented by the node embeddings 202. For example, the policy neural network 200 can have one or more tuning neural network layers that combine i) the outputs of previous neural network layers and ii) the node embeddings 202. As a specific example, the policy neural network 200 can use the following architecture to compute the output x of the dense neural network layer l. (l+1) :

[0059] x (l+1) =h (l) (c(x (0) )⊙x (l) ),

[0060] Where, x (l) It is the input of layer l in a dense neural network, g (l) The operation of layer l in a dense neural network is represented by c, which is a feature of the neural network layer in a policy-regulating neural network, and x (0) It is a node embedding 202.

[0061] Adjusting one or more layers of the policy neural network 200 according to the characteristics of the corresponding computation graph allows the policy neural network 200 to generate execution optimization policies for neural networks with very different architectures. For example, the policy neural network 200 can be trained to receive neural network graphs from a wide range of application areas (e.g., computer vision, natural language processing, and robotics). These neural networks can have drastically different architectures, different types of computational operations, data shapes, and network topologies, and can have vastly different numbers of operations ranging from, for example, 1K to 10M. Adjusting the neural network layers allows the policy neural network 200 to utilize information about a specific neural network to generate execution optimization policies for that neural network.

[0062] Figure 3 This is a diagram of example task subnet 300. Task subnet 300 is an example of a system implemented as a computer program on one or more computers in one or more locations, wherein the systems, components and technologies described below are capable of implementation.

[0063] Task subnet 300 receives a task embedding 302 corresponding to a specific execution optimization task and generates a task output 322 for that specific execution optimization task. Task subnet 300 can also generate a next task embedding 332 corresponding to different execution optimization tasks and which can be provided to different task subnets corresponding to different execution optimization tasks. Task subnet 300 can be, for example... Figure 2 Any of the task subnets 220a-c depicted in the diagram.

[0064] Task subnet 300 includes task attention subnet 310 and task policy subnet 320. Task attention subnet 310 processes task embedding 302 to generate attention embedding 312. Task attention subnet can include one or more attention neural network layers, such as attention neural network layers with segmented recursive mechanisms. Task policy subnet 320 processes attention embedding 312 to generate task output 322. Task policy subnet can include one or more fully connected neural network layers.

[0065] The task subnet 300 may also include a combining subnet that combines the task output 322 and the task embedding 302 to generate the next task embedding 332. In some embodiments, the combining subnet 330 may compute a sum of the task output 322 and the task embedding 302. In some such embodiments, the combining subnet may also use layer normalization to process the sum to generate the next task embedding 332. In some other embodiments, the combining subnet 330 may combine the task embedding 302 and the task output 322 by processing the task embedding 302 and the task output 322 using one or more neural network layers.

[0066] In some such implementations, with policy neural networks (e.g., Figure 2 In the policy neural network 200 depicted herein, each task attention subnetwork 310 corresponding to a different execution optimization task can share parameters. Parameter sharing allows the policy neural network to consider dependencies between multiple execution optimization tasks. As a specific example, each task attention subnetwork can include one or more of the same recurrent attention neural network layers. As another specific example, each task attention subnetwork can be the same subnetwork, i.e., having the same network architecture and parameter values. In some implementations, each task attention subnetwork 310 corresponding to a specific execution optimization task in the policy neural network can include residual connections to other task attention subnetworks corresponding to other execution optimization tasks. Residual connections allow the policy neural network to further capture inter-task dependencies.

[0067] Figure 4 This is a flowchart of an example process 400 for generating and implementing optimization strategies. For convenience, process 400 will be described as being executed by a system of one or more computers located in one or more locations. For example, a strategy generation system appropriately programmed according to this specification (e.g., Figure 1 The strategy generation system 100 described in the text is capable of executing process 400.

[0068] The system obtains data representing a graph characterizing the operations of the first neural network (step 402). The data can be represented by a graph representation subsystem (e.g., Figure 1 The graph representation subsystem 110 depicted in the figure is generated. The data can include a corresponding initial node embedding for each node in the graph, wherein each node in the graph represents a different operation of the first neural network.

[0069] The system uses a graph embedding neural network to process the data representing the graph to generate graph embeddings (step 404). Graph embeddings can include corresponding updated node embeddings generated at multiple embedding time steps for each node in the graph.

[0070] The system uses the embedding of the processing graph by a policy neural network to generate an execution optimization policy for the first neural network (step 406). The execution optimization policy can include one or more task outputs, where each task output corresponds to a specific execution optimization task, and includes a corresponding decision for the specific execution optimization task for each operation of the first neural network. The policy neural network can generate one or more task outputs by processing the graph embedding using multiple attention neural network layers with a piecewise recursive mechanism. In an implementation where the policy neural network generates multiple task outputs, the policy neural network can generate each task output sequentially using a sequence of task subnets.

[0071] The system provides the generated execution optimization strategy to one or more processing devices (step 408). The processing devices can use the generated execution optimization strategy to perform operations on the first neural network.

[0072] Figure 5 This is a flowchart of an example process 500 for generating multiple task outputs. Each task output corresponds to a different execution optimization task among multiple execution optimization tasks. For convenience, process 500 will be described as being executed by a system of one or more computers located in one or more locations. For example, a policy neural network appropriately programmed according to this specification (e.g., Figure 2 The policy neural network 200 described in the text is capable of executing process 500.

[0073] The system receives the graph embedding representing the operation of the first neural network (step 501). The graph embedding may have already been generated by a graph embedding neural network (e.g., Figure 1 The graph depicted in the diagram is generated by an embedded neural network 120.

[0074] At the first time step, the system processes the graph embedding to generate a first task embedding (step 502). The first time step and the first task embedding can correspond to a first execution optimization task among multiple execution optimization tasks. For example, the system can use a policy to embed a subnet (e.g., Figure 2 The strategy described in the text (embedding subnet 210) is used to handle graph embedding.

[0075] At the first time step, the system processes the first task embedding to generate a first task output (step 504). The first task output includes the corresponding decision for each operation of the first neural network for the first execution optimization task. For example, the system can use a task subnetwork corresponding to the first execution optimization task (e.g., Figure 2 The first task subnetwork 220a) of the policy neural network 200 described in the figure is used to process the first task embedding.

[0076] At a subsequent time step, the system processes the first task output of the first time step to generate a subsequent task embedding (step 506). The subsequent time step and the subsequent task embedding can correspond to a subsequent execution optimization task different from the first execution optimization task. For example, the system can use a combined subnet (e.g., Figure 3 The combined subnet 330 depicted in the diagram combines the first task output with the first task embedding to generate subsequent task embeddings.

[0077] At a subsequent time step, the system processes the subsequent task embedding to generate the subsequent task output (step 508). The subsequent task output includes the corresponding decisions for each operation of the first neural network for the subsequent execution of the optimization task.

[0078] The system determines whether all task outputs have been generated (step 510). That is, the system determines whether task outputs have been generated for each of the multiple execution optimization tasks, and thus determines whether the execution optimization strategy for the first neural network has been fully generated.

[0079] If not all task outputs are generated at step 510, the system returns to step 506. That is, at another subsequent time step corresponding to a different execution optimization task for which no task output has yet been generated, the system processes the task output of the previous time step to generate subsequent task embeddings for the different execution optimization tasks, and then processes the subsequent task outputs to generate subsequent task outputs for the different execution optimization tasks.

[0080] If all task outputs have been generated in step 510, the system provides the generated task outputs to one or more processing devices for performing the operations of the first neural network (step 512).

[0081] This specification uses the term "configuration" in conjunction with system and computer program components. For a system of one or more computers to be configured to perform a specific operation or action, it means that the system has software, firmware, hardware, or a combination thereof installed thereon, which, in operation, causes the system to perform the operation or action. For one or more computer programs to be configured to perform a specific operation or action, it means that one or more programs include instructions that, when executed by a data processing device, cause the device to perform the operation or action.

[0082] Embodiments of the subject matter and functional operation described in this specification can be implemented in digital electronic circuit systems, in tangibly embodied computer software or firmware, in computer hardware, including the structures disclosed in this specification and their structural equivalents, or combinations thereof. Embodiments of the subject matter described in this specification can be implemented as one or more computer programs, i.e., one or more modules of computer program instructions encoded on a tangible, non-transitory storage medium, for execution by a data processing device or for controlling the operation of a data processing device. The computer storage medium can be a machine-readable storage device, a machine-readable storage substrate, a random or serial access memory device, or combinations thereof. Alternatively or additionally, the program instructions can be encoded on artificially generated propagation signals (e.g., machine-generated electrical, optical, or electromagnetic signals), generated to encode information for transmission to a suitable receiver device for execution by the data processing device.

[0083] The term "data processing apparatus" refers to data processing hardware and includes all types of devices, apparatuses, and machines for processing data, including, for example, programmable processors, computers, or multiple processors or computers. The apparatus may also be or further include a dedicated logic circuit system, such as a FPGA (Field-Programmable Gate Array) or an ASIC (Application-Specific Integrated Circuit). In addition to hardware, the apparatus may optionally include code that creates an execution environment for computer programs, such as code constituting processor firmware, a protocol stack, a database management system, an operating system, or a combination of one or more of these.

[0084] A computer program, which may also be referred to or described as a program, software, software application, app, module, software module, script, or code, can be written in any form of programming language (including compiled or interpreted languages, or declarative or procedural languages); and it can be deployed in any form, including as a standalone program or as a module, component, subroutine, or other unit suitable for use in a computing environment. A program may, but does not necessarily, correspond to a file in a file system. A program can be stored as a portion of a file that holds other programs or data (e.g., one or more scripts stored in a markup language document), as a single file dedicated to the program in question, or as multiple coordinating files (e.g., a file storing one or more modules, subroutines, or code portions). A computer program can be deployed to execute on a single computer or on multiple computers located at a site or distributed across multiple sites and interconnected via a data communication network.

[0085] In this specification, the term "database" is used broadly to refer to any collection of data: data that does not need to be structured in any particular way, or not structured at all, and that can be stored on storage devices in one or more locations. Thus, for example, an indexed database can comprise multiple data collections, each of which can be organized and accessed differently.

[0086] Similarly, in this specification, the term "engine" is used broadly to refer to a software-based system, subsystem, or process programmed to perform one or more specific functions. Typically, an engine will be implemented as one or more software modules or components installed on one or more computers in one or more locations. In some cases, one or more computers will be dedicated to a particular engine; in others, multiple engines can be installed and run on the same one or more computers.

[0087] The processes and logic flows described in this specification can be executed by one or more programmable computers that execute one or more computer programs to perform functions by manipulating input data and generating output. The processes and logic flows can also be executed by special-purpose logic circuitry (e.g., FPGA or ASIC) or by a combination of special-purpose logic circuitry and one or more programmable computers.

[0088] A computer suitable for executing computer programs can be based on a general-purpose or special-purpose microprocessor, or both, or any other type of central processing unit. Typically, the central processing unit receives instructions and data from read-only memory or random access memory, or both. The basic components of a computer are the central processing unit for executing or running instructions and one or more memory devices for storing instructions and data. The central processing unit and memory can be supplemented by or incorporated into special-purpose logic circuitry. Typically, a computer will also include one or more mass storage devices for storing data, such as disks, magneto-optical disks, or optical disks, or operatively coupled to receive data from or transfer data to them, or both. However, a computer does not need to have such devices. Furthermore, a computer can be embedded in another device (e.g., a mobile phone, a personal digital assistant (PDA), a mobile audio or video player, a game console, a global positioning system (GPS) receiver, or a portable storage device such as a universal serial bus (USB) flash drive, to name a few).

[0089] Computer-readable media suitable for storing computer program instructions and data include all forms of non-volatile memory, media, and memory devices, including, for example, semiconductor memory devices such as EPROM, EEPROM, and flash memory devices; magnetic disks such as internal hard disks or removable disks; magneto-optical disks; and CD ROMs and DVD-ROMs.

[0090] To provide interaction with the user, embodiments of the subject matter described in this specification can be implemented on a computer having a display device for displaying information to the user (e.g., a CRT (cathode ray tube) or LCD (liquid crystal display) monitor) and a keyboard and pointing device (e.g., a mouse or trackball) through which the user can provide input to the computer. Other types of devices can also be used to provide interaction with the user; for example, feedback provided to the user can be any form of sensory feedback, such as visual, auditory, or tactile feedback; and input from the user can be received in any form, including sound, speech, or tactile input. Additionally, the computer can interact with the user by sending documents to and receiving documents from the device used by the user; for example, by sending a webpage to a web browser on the user's device in response to a request received from a web browser. Furthermore, the computer can interact with the user by sending text messages or other forms of messages to a personal device (e.g., a smartphone running a messaging application) and receiving response messages in return from the user.

[0091] The data processing apparatus used to implement machine learning models may also include, for example, dedicated hardware accelerator units for processing the common and computationally intensive parts of machine learning training or production, namely inference and workloads.

[0092] It is possible to implement and deploy machine learning models using machine learning frameworks (such as TensorFlow, Microsoft Cognitive Toolkit, Apache Singa, or Apache MXNet).

[0093] Embodiments of the subject matter described in this specification can be implemented in a computing system that includes back-end components (e.g., as a data server), or middleware components (e.g., an application server), or front-end components (e.g., a client computer having a graphical user interface, web browser, or app that a user can interact with through embodiments of the subject matter described in this specification), or any combination of one or more such back-end, middleware, or front-end components. The components of the system can be interconnected via digital data communication (e.g., a communication network) of any form or medium. Examples of communication networks include local area networks (LANs) and wide area networks (WANs), such as the Internet.

[0094] A computing system can include clients and servers. Clients and servers are typically geographically separated and usually interact via a communication network. The client-server relationship is established by means of computer programs running on respective computers and having a client-server relationship with each other. In some embodiments, the server transmits data (e.g., HTML pages) to a user device, for example, for the purpose of displaying data to a user interacting with the device acting as a client and receiving user input from it. It is possible to receive data generated at the user device, such as the result of user interaction, from the device at the server.

[0095] In addition to the embodiments described above, the following embodiments are also innovative:

[0096] Example 1 is a method for generating task outputs for performing multiple operations of a neural network on one or more processing devices, wherein: the task outputs include corresponding decisions for each of the multiple operations of the neural network for a specific optimization task, the method comprising:

[0097] Obtain data representing a graph that characterizes multiple operations of a neural network, where each node of the graph represents an operation of the neural network and each edge of the graph represents a data dependency between operations;

[0098] Graph embedding neural networks are used to process data representing graphs to generate graph embeddings; and

[0099] Use a policy neural network to process the embeddings of the graph to generate task output.

[0100] Example 2 is the method according to Example 1, wherein the graph embedding includes the corresponding node embedding of each node of the graph.

[0101] Example 3 is the method according to Example 2, wherein processing the data representing the graph using a graph embedding neural network includes at each of a plurality of embedding time steps:

[0102] Receive the current embedding of each node in the graph generated during the previous embedding time step;

[0103] For each specific node in the graph, combine the corresponding current embeddings of each of the node's neighboring nodes to generate the node's neighborhood embedding; and

[0104] For each node in the graph, the node's current embedding and the node's neighborhood embeddings are combined to generate a new embedding for the node.

[0105] Example 4 is the method according to Example 3, wherein processing the data representing the graph using a graph embedding neural network further includes, at the first embedding time step:

[0106] For each node in the graph, an initial embedding is generated using the node's features, where the features include one or more of the following:

[0107] The operation type represented by the node.

[0108] The output shape of the operation represented by the node, or

[0109] A corresponding identifier for each of the node's neighboring nodes.

[0110] Example 5 is the method according to any one of Examples 3 or 4, wherein combining the current embedding of a particular node and the neighborhood embedding of the particular node includes:

[0111] Cascade the current embedding of a specific node and the neighborhood embedding of the specific node to generate a combined embedding of the specific node; and

[0112] Use one or more fully connected neural network layers to process combined embeddings of a specific node to generate a new embedding for that specific node.

[0113] Example 6 is the method according to any one of Examples 3 to 5, wherein combining the corresponding embeddings of each neighboring node of a particular node includes:

[0114] For each neighboring node of a given node, an affine transformation is used to process the current embedding of the neighboring nodes to generate a processed embedding for the neighboring nodes.

[0115] For each neighboring node of a specific node, the sigmoid activation function is used to process the processed embeddings of the neighboring nodes to generate activation embeddings for the neighboring nodes; and

[0116] The activation embeddings of each neighboring node of the given node are combined by using a max-pooling layer to process the activation embeddings.

[0117] Example 7 is the method according to any one of Examples 1 to 6, wherein the specific optimization task is one or more of the following:

[0118] Device placement task, in which each of the multiple operations of the neural network is assigned to a specific processing device in one or more processing devices;

[0119] An operation scheduling task is performed, wherein each of a plurality of operations of a neural network is assigned a priority, and wherein each processing device includes a priority-based scheduler that maintains a priority queue of operations assigned to the processing device; or

[0120] Operation fusion task, in which multiple selected operations of a neural network are determined to be executed as a single operation.

[0121] Example 8 is the method according to any one of Examples 1 to 7, wherein a graph embedding neural network and a policy neural network have been trained end-to-end by updating the parameters θ of the neural network using an objective function, which characterizes the expected runtime of a corresponding operation represented by each of a plurality of candidate graphs.

[0122] Example 9 is the method described according to Example 8, wherein the objective function is:

[0123]

[0124] Where G is a candidate graph. It is the space of candidate graphs, T is the task output for a specific optimization task, and π θ (G) is the strategy of the candidate graph G under parameter θ, and r G,T The reward is a function of the runtime of the operation represented by the candidate graph G using the task output T.

[0125] Example 10 is the method according to any one of Examples 8 to 9, wherein the graph embedding neural network and the policy neural network have been trained by optimizing the objective function using a neighbor policy.

[0126] Example 11 is the method according to any one of Examples 1 to 10, wherein the policy neural network is adjusted according to the features of the graph.

[0127] Example 12 is the method according to Example 11, wherein adjusting the policy neural network based on the features of the graph includes calculating the output x of the neural network layer l of the policy neural network. (l+1) :

[0128] x (l+1) =g (l) (c(x (0) )⊙x (l) ),

[0129] Where, x (l) It is the input of neural network layer l, g (l) The operation of the neural network layer l is characterized by c, which is a feature of the neural network layer of the regulation policy neural network, and x (0) It is an embedding of a graph generated by a graph embedding neural network.

[0130] Example 13 is a method according to any one of Examples 1 to 12, wherein the method further includes performing multiple operations of a neural network on one or more processing devices using the generated task output.

[0131] Example 14 is the method according to any one of Examples 1 to 13, wherein the policy neural network includes multiple attention neural network layers.

[0132] Example 15 is the method according to Example 14, wherein the multiple attention neural network layers include a segmented recursive mechanism.

[0133] Example 16 is a method according to any one of Examples 1 to 15, wherein the method further includes generating a corresponding task output for each of a plurality of optimization tasks, wherein each task output includes a corresponding decision for each of a plurality of operations of the neural network for the corresponding optimization task, and wherein using a policy neural network to process the graph embedding includes:

[0134] At the first time step corresponding to the first optimization task in a plurality of optimization tasks:

[0135] Generate the first task embedding from the graph embedding; and

[0136] A policy neural network is used to process the first task embedding to generate a first task output for a first optimization task; and

[0137] At each of one or more subsequent time steps, each corresponding to a specific optimization task among multiple optimization tasks:

[0138] Process the task output of the previous time step to generate the embedding of the subsequent task; and

[0139] The subsequent task embedding is processed using a policy neural network to generate subsequent task outputs for the corresponding optimized task.

[0140] Example 17 is the method according to Example 16, wherein processing the task output of the previous time step includes combining the task embedding of the previous time step and the task output of the previous time step to generate the subsequent task embedding.

[0141] Example 18 is the method according to Example 17, wherein combining the task embedding of the previous time step and the task output of the previous time step to generate the subsequent task embedding includes:

[0142] Calculate the sum of the task embeddings and task outputs from the previous time step to generate a summation representation; and

[0143] Layer normalization is used to process the summation representation to generate subsequent task embeddings.

[0144] Example 19 is the method according to any one of Examples 16 to 18, wherein:

[0145] The policy neural network includes a policy embedding subnetwork, which comprises multiple attention neural network layers; and

[0146] Generating the first task embedding from the graph embedding involves using a policy embedding subnet to process the graph embedding.

[0147] Example 20 is the method according to any one of Examples 16 to 19, wherein the policy neural network comprises:

[0148] A corresponding task attention subnetwork for each of multiple optimization tasks, wherein each task attention subnetwork includes one or more attention neural network layers; and

[0149] A corresponding task policy subnet for each of the multiple optimization tasks.

[0150] Among these, using policy neural networks to process task embeddings for specific optimization tasks includes:

[0151] Task embeddings are processed using a task attention subnet corresponding to a specific optimization task to generate attention embeddings; and

[0152] Attention embeddings are processed using a task policy subnet corresponding to a specific optimization task to generate task outputs for that specific optimization task.

[0153] Example 21 is the method according to Example 20, wherein the task attention subnet of each of the multiple optimization tasks shares multiple parameters.

[0154] Example 22 is the method according to any one of Examples 20 to 21, wherein the corresponding task attention subnet for each optimization task shares one or more residual connections with each of one or more other task attention subnets for other optimization tasks.

[0155] Example 23 is a method for generating an optimized strategy for performing multiple operations of a neural network on one or more processing devices, wherein:

[0156] The optimization strategy includes the task output for each of the multiple optimization tasks, and

[0157] For each of the multiple operations of the neural network, the output of each task includes a corresponding decision for the corresponding optimization task.

[0158] The method includes:

[0159] Obtain data representing a graph that characterizes multiple operations of a neural network, where each node of the graph represents an operation of the neural network and each edge of the graph represents a data dependency between operations;

[0160] Graph embedding neural networks are used to process data representing graphs to generate graph embeddings; and

[0161] The embedding of the graph is processed using a policy neural network to generate a corresponding task output for each of a plurality of optimization tasks, the processing including:

[0162] At the first time step corresponding to the first optimization task in a plurality of optimization tasks:

[0163] Generate the first task embedding from the graph embedding; and

[0164] A policy neural network is used to process the first task embedding to generate a first task output for a first optimization task; and

[0165] At each of one or more subsequent time steps, each corresponding to a specific optimization task among multiple optimization tasks:

[0166] Process the task output of the previous time step to generate the embedding of the subsequent task; and

[0167] A policy neural network is used to process the subsequent task embeddings to generate subsequent task outputs for the corresponding optimized task.

[0168] Example 24 is the method according to Example 23, wherein processing the task output of the previous time step includes combining the task embedding of the previous time step and the task output of the previous time step to generate the subsequent task embedding.

[0169] Example 25 is the method according to Example 24, wherein combining the task embedding of the previous time step and the task output of the previous time step to generate the subsequent task embedding includes:

[0170] Calculate the sum of the task embeddings and task outputs from the previous time step to generate a summation representation; and

[0171] Layer normalization is used to process the summation representation to generate subsequent task embeddings.

[0172] Example 26 is the method according to any one of Examples 23 to 25, wherein the graph embedding includes the corresponding node embedding of each node of the graph.

[0173] Example 27 is the method according to Example 26, wherein using a graph embedding neural network to process data representing a graph includes at each of a plurality of embedding time steps:

[0174] Receive the current embedding of each node in the graph generated during the previous embedding time step;

[0175] For each specific node in the graph, combine the corresponding current embeddings of each of the node's neighboring nodes to generate the node's neighborhood embedding; and

[0176] For each node in the graph, the node's current embedding and the node's neighborhood embeddings are combined to generate a new embedding for the node.

[0177] Example 28 is the method according to Example 27, wherein processing the data representing the graph using a graph embedding neural network further includes, at the first embedding time step:

[0178] For each node in the graph, an initial embedding is generated using the node's features, where the features include one or more of the following:

[0179] The operation type represented by the node.

[0180] The output shape of the operation represented by the node, or

[0181] A corresponding identifier for each of the node's neighboring nodes.

[0182] Example 29 is the method according to any of Examples 27 to 28, wherein combining the current embedding of a particular node and the neighborhood embedding of the particular node includes:

[0183] Cascade the current embedding of a specific node and the neighborhood embedding of the specific node to generate a combined embedding of the specific node; and

[0184] Use one or more fully connected neural network layers to process combined embeddings of a specific node to generate a new embedding for that specific node.

[0185] Example 30 is the method according to any one of Examples 27 to 29, wherein combining the corresponding embeddings of each neighboring node of a particular node includes:

[0186] For each neighboring node of a given node, an affine transformation is used to process the current embedding of the neighboring nodes to generate a processed embedding for the neighboring nodes.

[0187] For each neighboring node of a specific node, the sigmoid activation function is used to process the processed embeddings of the neighboring nodes to generate activation embeddings for the neighboring nodes; and

[0188] The activation embeddings of each neighboring node of a particular node are combined by using a max-pooling layer to process the activation embeddings.

[0189] Example 31 is the method according to any one of Examples 23 to 30, wherein the plurality of optimization tasks include one or more of the following:

[0190] Device placement task, in which each of the multiple operations of the neural network is assigned to a specific processing device in one or more processing devices;

[0191] An operation scheduling task is performed, wherein each of a plurality of operations of a neural network is assigned a priority, and wherein each processing device includes a priority-based scheduler that maintains a priority queue of operations assigned to the processing device; or

[0192] Operation fusion task, in which multiple selected operations of a neural network are determined to be executed as a single operation.

[0193] Example 32 is the method according to any one of Examples 23 to 31, wherein a graph embedding neural network and a policy neural network have been trained end-to-end by updating the parameters θ of the neural network using an objective function, which characterizes the expected runtime of a corresponding operation represented by each of a plurality of candidate graphs.

[0194] Example 33 is the method according to Example 32, wherein the objective function is:

[0195]

[0196] Where G is a candidate graph. It is the space of candidate graphs, and T is the specific optimization task used in multiple optimization tasks. Task output, π θ (G) is the strategy of the candidate graph G under parameter θ, and r G,T The reward is a function of the runtime of the operation represented by the candidate graph G using the task output T.

[0197] Example 34 is the method according to any of Examples 32 to 33, wherein the graph embedding neural network and the policy neural network have been trained by optimizing the objective function using a neighbor policy.

[0198] Example 35 is the method according to any one of Examples 23 to 34, wherein the policy neural network is adjusted according to the features of the graph.

[0199] Example 36 is the method according to Example 35, wherein adjusting the policy neural network based on the features of the graph includes calculating the output x of the neural network layer l of the policy neural network. (l+1) :

[0200] x (l+1) =g (l) (c(x (0) )⊙x (l) ),

[0201] Where, x (l) It is the input of the neural network layer l, g (l) The operation of the neural network layer l is characterized by c, which is a feature of the neural network layer of the regulation policy neural network, and x (0) It is an embedding of a graph generated by a graph embedding neural network.

[0202] Example 37 is the method according to any one of Examples 23 to 36, further comprising:

[0203] The generated optimization strategy is used to perform multiple operations of the neural network on one or more processing devices.

[0204] Example 38 is the method according to any one of Examples 23 to 37, wherein:

[0205] The policy neural network includes a policy embedding subnetwork, which comprises multiple recurrent attention neural network layers; and

[0206] Generating the first task embedding from the graph embedding involves using a policy embedding subnet to process the graph embedding.

[0207] Example 39 is the method according to any one of Examples 23 to 38, wherein the policy neural network includes:

[0208] A corresponding task attention subnetwork for each of multiple optimization tasks, wherein each task attention subnetwork includes one or more attention neural network layers; and

[0209] A corresponding task policy subnet for each of the multiple optimization tasks.

[0210] Among them, the use of policy neural networks to process task embeddings for specific optimization tasks includes:

[0211] Task embeddings are processed using a task attention subnet corresponding to a specific optimization task to generate attention embeddings; and

[0212] Attention embeddings are processed using a task policy subnet corresponding to a specific optimization task to generate task outputs for that specific optimization task.

[0213] Example 40 is the method according to Example 39, wherein the task attention subnet of each of the multiple optimization tasks shares multiple parameters.

[0214] Example 41 is the method according to any one of Examples 39 to 40, wherein the corresponding task attention subnet for each optimization task shares one or more residual connections with each of one or more other task attention subnets for other optimization tasks.

[0215] Example 42 is a system comprising: a system including one or more computers and one or more storage devices storing instructions, the instructions being operable, when executed by the one or more computers, to cause the one or more computers to perform the method according to any one of Examples 1 to 41.

[0216] Example 43 is a computer storage medium encoded with a computer program, the program including instructions that, when executed by a data processing device, are operable to cause the data processing device to perform the method according to any one of Examples 1 to 41.

[0217] While this specification contains numerous details of specific implementations, these should not be construed as limiting the scope of any invention or the scope that may be claimed, but rather as descriptions of features that may be specific to particular embodiments of a particular invention. Some features described in the context of individual embodiments in this specification can also be implemented in combination in a single embodiment. Conversely, various features described in the context of a single embodiment can also be implemented individually or in any suitable sub-combination in multiple embodiments. Furthermore, although features may be described above as functioning in certain combinations and even initially claimed in this way, in some cases it is possible to remove one or more features from the claimed combination, and the claimed combination may involve sub-combinations or variations of sub-combinations.

[0218] Similarly, although operations are depicted in the accompanying drawings and described in a specific order in the claims, this should not be construed as requiring these operations to be performed in the specific order shown or sequentially, or to perform all of the shown operations to achieve the desired result. In some cases, multitasking and parallel processing may be advantageous. Furthermore, the separation of various system modules and components in the above embodiments should not be construed as requiring such separation in all embodiments, and it should be understood that the described program components and systems can generally be integrated together in a single software product or packaged into multiple software products.

[0219] Specific embodiments of the subject matter have been described. Other embodiments are within the scope of the following claims. For example, the actions recited in the claims can be performed in different orders and still achieve the desired result. As an example, the processes depicted in the drawings do not necessarily require the specific order or sequential order shown to achieve the desired result. In some cases, multitasking and parallel processing may be advantageous.

Claims

1. A method for generating an execution optimization strategy for performing multiple operations of a neural network on one or more processing devices, wherein: The execution optimization strategy includes task outputs for each of a plurality of execution optimization tasks, wherein the plurality of execution optimization tasks include one or more of the following: Device placement task, wherein each of the plurality of operations of the neural network is assigned to a specific processing device among the one or more processing devices. The operation scheduling task, wherein each of the plurality of operations of the neural network is assigned a priority, and wherein each processing device includes a priority-based scheduler that maintains a priority queue of the operations assigned to the processing device, or An operation fusion task is performed, wherein multiple selected operations of the neural network are merged such that a specific processing device assigned to execute the selected operations is configured to treat the selected operations as a single operation, and wherein the task output for this execution optimization task is the identifiers of two or more nodes of the computation graph to be treated as a single operation by the assigned processing device, and For each of the plurality of operations of the neural network, each task output includes a corresponding decision for performing the optimization task. The method includes: Obtain data representing a graph characterizing the plurality of operations of the neural network, wherein each node of the graph characterizes an operation of the neural network and each edge of the graph characterizes a data dependency between the operations; A graph embedding neural network is used to process the data representing the graph to generate embeddings of the graph, wherein the embeddings are ordered sets of numerical values ​​representing inputs in a specific embedding space, wherein the numerical values ​​have a fixed dimension; and The embedding of the graph is processed using a policy neural network to generate a corresponding task output for each of the plurality of execution optimization tasks, the processing including: At the first time step corresponding to the first execution optimization task among the plurality of execution optimization tasks: A first task embedding is generated from the embedding of the graph; and The policy neural network is used to process the first task embedding to generate a first task output for the first execution optimization task; and At each of one or more subsequent time steps corresponding to the respective execution optimization task among the plurality of execution optimization tasks: Process the task output from the previous time step to generate the subsequent task embedding; and The policy neural network is used to process the subsequent task embedding to generate subsequent task outputs for corresponding execution optimization tasks.

2. The method according to claim 1, wherein, Processing the task output of a previous time step includes combining the task embedding of the previous time step and the task output of the previous time step to generate the subsequent task embedding.

3. The method according to claim 2, wherein, Combining the task embedding from the previous time step and the task output from the previous time step to generate the subsequent task embedding includes: Calculate the sum of the task embedding and the task output at the previous time step to generate a sum representation; and Layer normalization is used to process the summation representation to generate the subsequent task embedding.

4. The method according to claim 1, wherein, The embedding of the graph includes the corresponding node embedding for each node of the graph.

5. The method according to claim 4, wherein, The graph embedding neural network is used to process the data representing the graph, including at each of the multiple embedding time steps: Receive the current embedding of each node in the graph generated during the previous embedding time step; For each specific node in the graph, the corresponding current embeddings of each neighboring node of the specific node are combined to generate the neighborhood embedding of the specific node; as well as For each node in the graph, the current embedding of the node and the neighborhood embedding of the node are combined to generate a new embedding for the node.

6. The method according to claim 5, wherein, Using the graph embedding neural network to process the data representing the graph further includes, at the first embedding time step: An initial embedding is generated for each node of the graph using the features of the node, wherein the features include one or more of the following: The operation type represented by the node. The output shape of the operation's output, characterized by the node, and A corresponding identifier for each of the node's neighboring nodes.

7. The method according to claim 5, wherein, Combining the current embedding of a specific node and the neighborhood embedding of the specific node includes: Cascade the current embedding of the specific node and the neighborhood embedding of the specific node to generate a combined embedding of the specific node; and One or more fully connected neural network layers are used to process the combined embeddings of the particular node to generate the new embedding of the particular node.

8. The method according to claim 5, wherein, The corresponding embeddings of each neighboring node of a specific node include: For each neighboring node of the given node, an affine transformation is used to process the current embedding of the neighboring node to generate a processed embedding for the neighboring node; For each neighboring node of the specific node, the sigmoid activation function is used to process the processed embeddings of the neighboring nodes to generate activation embeddings for the neighboring nodes; and The activation embeddings of each neighboring node of the particular node are combined by processing the activation embeddings using a max-pooling layer.

9. The method according to claim 1, wherein, The graph embedding neural network and the policy neural network have been trained end-to-end by updating the parameters θ of the neural network using an objective function, which represents the expected runtime of the corresponding operation represented by each of the multiple candidate graphs.

10. The method according to claim 9, wherein, The objective function is: Where G is a candidate graph. It is the space of candidate graphs, and T is the specific execution optimization task among the plurality of execution optimization tasks. Task output, π θ (G) is the strategy of the candidate graph G under the parameter θ, and r G,T The reward is a function of the running time of the operation represented by the candidate graph G using the task output T.

11. The method according to claim 9, wherein, The graph embedding neural network and the policy neural network have been trained by optimizing the objective function using a proximity policy.

12. The method according to claim 1, wherein, The policy neural network is adjusted based on the features of the graph.

13. The method according to claim 12, wherein, Adjusting the policy neural network based on the features of the graph includes calculating the output x of the neural network layer l of the policy neural network. (l+1) : x (l+1) =g (l) (c(x (0) )⊙x (l) ), Where, x (l) It is the input of the neural network layer l, g (l) The operation of the neural network layer l is characterized by c, which is a feature of the neural network layer that regulates the policy neural network, and x (0) The embedding of the graph is generated by the graph embedding neural network.

14. The method of claim 1, further comprising: The generated optimization strategy is used to perform the plurality of operations of the neural network on one or more processing devices.

15. The method according to claim 1, wherein: The policy neural network includes a policy embedding subnetwork, which includes multiple recurrent attention neural network layers. as well as Generating the first task embedding from the embedding of the graph includes processing the embedding of the graph using the policy embedding subnet.

16. The method according to any one of claims 1 to 15, wherein, The policy neural network includes: A corresponding task attention subnetwork for each of the plurality of execution optimization tasks, wherein each task attention subnetwork includes one or more attention neural network layers; and A corresponding task policy subnet for each of the plurality of execution optimization tasks. The use of the policy neural network to process task embeddings for specific execution optimization tasks includes: The task embedding is processed using the task attention subnet corresponding to the specific execution optimization task to generate an attention embedding; and The attention embedding is processed using the task policy subnet corresponding to the specific execution optimization task to generate the task output for the specific execution optimization task.

17. The method according to claim 16, wherein, The task attention subnet of each of the plurality of execution optimization tasks shares multiple parameters.

18. The method according to claim 16, wherein, The corresponding task attention subnet for each optimization task shares one or more residual connections with each of one or more other task attention subnets for other optimization tasks.

19. A system comprising one or more computers and one or more storage devices storing instructions, the instructions being operable, when executed by the one or more computers, to cause the one or more computers to perform the method according to any one of claims 1 to 18.

20. A computer storage medium encoded with a computer program, the program including instructions that, when executed by a data processing apparatus, are operable to cause the data processing apparatus to perform the method according to any one of claims 1 to 18.

Images