Training neural networks represented as computational graphs
Systems and Methods for training a neural network represented as a computational graph are disclosed. An example method begins with obtaining data representing a computational graph. The computational graph is then augmented to generate a training computational graph for training the neural network using a machine learning training algorithm that includes computing a gradient of an objective function with respect to each of the parameters of the neural network. Augmenting the computational graph includes inserting a plurality of gradient nodes and training edges into the computational graph to generate a backward path through the computational graph that represents operations for computing the gradients of the objective function with respect to the parameters of the neural network. The neural network is trained using the machine learning training algorithm by executing the training computational graph.
1 . A method for training a neural network represented as a computational graph, wherein the computational graph comprises a plurality of nodes and a plurality of directed edges, the method comprising:
obtaining data representing the computational graph;
augmenting the computational graph to generate a training computational graph for training the neural network using a machine learning training algorithm that includes computing a gradient of an objective function with respect to each of one or more parameters of the neural network, comprising:
inserting a plurality of gradient nodes and training edges into the computational graph to generate a backward path through the computational graph that represents operations for computing the gradients of the objective function with respect to the one of more parameters of the neural network, wherein each gradient node represents a gradient function that computes a gradient of the objective function with respect to parameters flowing along a respective parameter directed edge in the computational graph;
in response to identifying a control flow node in the computational graph that causes one or more particular nodes in the computational graph to be iteratively performed during execution of the computational graph, augmenting the computational graph by inserting into the computational graph a monitoring node that (i) monitors a number of iterations of the one or more particular nodes performed during execution of the computational graph, and (ii) for each of multiple performed iterations of the one or more particular nodes, stores a tensor output from the one or more particular nodes for that iteration such that the monitoring node stores multiple tensors output from the one or more particular nodes corresponding to the multiple performed iterations, wherein, during execution of the backward path in the training computational graph, the multiple tensors stored by the monitoring node are provided as inputs to the gradient functions represented by one or more of the gradient nodes; and
training the neural network using the machine learning training algorithm by executing the training computational graph.
2 . The method of claim 1 , wherein augmenting the computational graph further comprises:
identifying one or more control flow nodes in the computational graph, wherein each control flow node represents a control flow operation that causes operations represented by one or more other nodes in the computational graph to be conditionally or iteratively performed during execution of the computational graph; and
inserting, for each identified control flow node, a corresponding backward path control flow node along the backward path through the computational graph, wherein each backward path control flow node adjusts control flow along the backward path to account for the control flow operation represented by the corresponding control flow node.
3 . The method of claim 2 , wherein the one or more control flow nodes include a switch control flow node, and wherein the backward path control flow node corresponding to the switch control flow node is a merge backward control flow node.
4 . The method of claim 2 , wherein the one or more control flow nodes include a merge control flow node, and wherein the backward path control flow node corresponding to the merge control flow node is a switch backward control flow node.
5 . The method of claim 2 , wherein the one or more control flow nodes include an enter control flow node, and wherein the backward path control flow node corresponding to the enter control flow node is an exit backward control flow node.
6 . The method of claim 2 , wherein the one or more control flow nodes include an exit control flow node, and wherein the backward path control flow node corresponding to the exit control flow node is an enter backward control flow node.
7 . The method of claim 2 , wherein the one or more control flow nodes include an iteration counter control flow node, and wherein the backward path control flow node corresponding to the iteration counter control flow node is an iteration counter backward control flow node.
8 . The method of claim 1 , further comprising:
determining that multiple iterations of the one or more particular nodes in the computational graph are performed during execution of the computational graph by analyzing the computational graph to identify the control flow node in the computational graph that cause the one or more particular nodes in the computational graph to be performed multiple times; and
wherein the monitoring node is inserted into the computational graph by automatically rewriting the computational graph.
9 . The method of claim 1 , wherein the neural network is a recurrent neural network that receives a respective neural network input at each of a plurality of time steps and generates a respective neural network output at each of the plurality of time steps, wherein the one or more particular nodes generate a respective node output for each of the plurality of time steps, and wherein the monitoring node stores the respective node outputs for each of the plurality of time steps.
10 . The method of claim 1 , wherein storing the tensor output from the one or more particular nodes includes:
asynchronously sending data representing the tensor output from the one or more particular nodes for an iteration from a device on which the tensor was produced to a central processing unit for storage after the data representing the tensor was produced; and
asynchronously retrieving the data representing the tensor from the central processing unit for use on the device in the backward path through the computational graph that represents operations for computing the gradient of the objective function with respect to the parameters of the neural network.
11 . The method of claim 10 , wherein the device is a graphical processing unit.
12 . The method of claim 1 , wherein training the neural network using the machine learning training algorithm by executing the training computational graph comprises:
allocating the nodes in the training computational graph across a plurality of devices; and
causing each of the devices to perform the operations represented by the nodes allocated to the device.
13 . The method of claim 1 , wherein the computational graph further comprises a plurality of connector directed edges, and a plurality of parameter directed edges, wherein each node represents a respective operation performed by the neural network as part of determining a neural network output from a neural network input, wherein each connector directed edge connects a respective first node to a respective second node that represents an operation that receives, as input, an output of an operation represented by the respective first node, and wherein each parameter directed edge connects into a respective node and represents a flow of one or more parameters of the neural network as input to the operation represented by the respective node.
14 . A system for training a neural network represented as a computational graph, wherein the computational graph comprises a plurality of nodes and a plurality of directed edges, the system comprising:
one or more computers; and
one or more storage devices storing instructions that are operable, when executed by the one or more computers, to cause the one or more computers to perform operations comprising:
obtaining data representing the computational graph;
augmenting the computational graph to generate a training computational graph for training the neural network using a machine learning training algorithm that includes computing a gradient of an objective function with respect to each of one or more parameters of the neural network, comprising:
inserting a plurality of gradient nodes and training edges into the computational graph to generate a backward path through the computational graph that represents operations for computing the gradients of the objective function with respect to the one or more parameters of the neural network, wherein each gradient node represents a gradient function that computes a gradient of the objective function with respect to parameters flowing along a respective parameter directed edge in the computational graph;
in response to identifying a control flow node in the computational graph that causes one or more particular nodes in the computational graph to be iteratively performed during execution of the computational graph, augmenting the computational graph by inserting into the computational graph a monitoring node that (i) monitors a number of iterations of the one or more particular nodes performed during execution of the computational graph, and (ii) for each of multiple performed iterations of the one or more particular nodes, stores a tensor output from the one or more particular nodes for that iteration such that the monitoring node stores multiple tensors output from the one or more particular nodes corresponding to the multiple performed iterations, wherein, during execution of the backward path in the training computational graph, the multiple tensors stored by the monitoring node are provided as inputs to the gradient functions represented by one or more of the gradient nodes; and
training the neural network using the machine learning training algorithm by executing the training computational graph.
15 . The system of claim 14 , wherein augmenting the computational graph further comprises:
identifying one or more control flow nodes in the computational graph, wherein each control flow node represents a control flow operation that causes operations represented by one or more other nodes in the computational graph to be conditionally or iteratively performed during execution of the computational graph; and
inserting, for each identified control flow node, a corresponding backward path control flow node along the backward path through the computational graph, wherein each backward path control flow node adjusts control flow along the backward path to account for the control flow operation represented by the corresponding control flow node.
16 . The system of claim 15 , wherein the one or more control flow nodes include a switch control flow node, and wherein the backward path control flow node corresponding to the switch control flow node is a merge backward control flow node.
17 . The system of claim 15 , wherein the one or more control flow nodes include a merge control flow node, and wherein the backward path control flow node corresponding to the merge control flow node is a switch backward control flow node.
18 . The system of claim 15 , wherein the one or more control flow nodes include an enter control flow node, and wherein the backward path control flow node corresponding to the enter control flow node is an exit backward control flow node.
19 . The system of claim 15 , wherein the one or more control flow nodes include an exit control flow node, and wherein the backward path control flow node corresponding to the exit control flow node is an enter backward control flow node.
20 . The system of claim 15 , wherein the one or more control flow nodes includes an iteration counter control flow node, and wherein the backward path control flow node corresponding to the iteration counter control flow nodes is an iteration counter backward control flow node.