pytorch-tree-lstm 0.1.1

A Tree-LSTM model package for PyTorch

pytorch-tree-lstm

This repo contains a PyTorch implementation of the child-sum Tree-LSTM model (Tai et al. 2015) implemented with vectorized tree evaluation and batching. This module has been tested with Python 3.6.6, PyTorch 0.4.0, and PyTorch 1.0.1.

High-level Approach

Efficient batching of tree data is complicated by the need to have evaluated all of a node's children before we can evaluate the node itself. To minimize the performance impact of this issue, we break the node evaluation process into steps such that at each step we evaluate all nodes for which all child nodes have been previously evaluated. This allows us to evaluate multiple nodes with each torch operation, increasing computation speeds by an order of magnitude over recursive approaches.

As an example, consider the following tree:

On the first step of the tree calculation, we can evaluate nodes 1 & 3 in parallel as neither has any child nodes. At the second step we are able to evaluate node 2, as its child node 3 was evaluated previously. Lastly we evaluate node 0, which depends on nodes 1 and 2. Doing this we can reduce a four-node computation to three steps. Bigger trees with more leaf nodes will experience larger performance gains.

To facilitate this approach we encode the Tree structure and features into four Tensors. For a tree with N nodes, E edges, and F features, the required Tensors are:

• features - A size N x F tensor containing the features for each node.
• adjacency_list - A size E x 2 tensor containing the node indexes of the parent node and child node for every connection in the tree.
• node_order - A size N tensor containing the calculation step at which a node can be evaluated. Note that the order that node data is stored in features and node_order must be identical.
• edge_order - A size E tensor containing the calculation step at which each entry in the adjacency_list is needed in order to retrieve the child nodes for a current node. Note that the order that parent-child data is stored in adjacency_list and edge_order must be identical.

node_order and edge_order hold redundant information derivable from the adjacency_list and features; however, precomputing these tensors gives a significant performance improvement due to the current lack of an efficient set intersection function in PyTorch 1.0. The order tensors can be generated using the treelstm.calculate_evaluation_orders function. calculate_evaluation_orders accepts the adjacency_list tensor and the length of the features tensor and returns the two order tensors:

import treelstm
node_order, edge_order = treelstm.calculate_evaluation_orders(adjacency_list, len(features))

The tensor representation of the example tree above would be:

features: tensor([[1., 0.],
                  [0., 1.],
                  [0., 0.],
                  [1., 1.]])

adjacency_list: tensor([[0, 1],
                        [0, 2],
                        [2, 3]])

node_order: tensor([2, 0, 1, 0])

edge_order: tensor([2, 2, 1])

Installation

The pytorch-tree-lstm package can be installed via pip:

pip install pytorch-tree-lstm

Once installed, the library can be imported via:

import treelstm

Usage

The file tree_list.py contains the TreeLSTM module. The module accepts the features, node_order, adjacency_list, edge_order tensors detailed above as input.

These tensors can be batched together by concatenation (torch.cat()) with the exception of the adjacency_list. The adjacency_list contains indexes into the features tensor used to retrieve child features for performing sums over node children, and when batched together these indexes must be adjusted for the new position of the features in the batched tensors.

The treelstm.batch_tree_input function is provided to do this concatenation and adjustment. treelstm.batch_tree_input accepts a list of dictionaries containing fields features, node_order, adjacency_list, and edge_order and returns a dictionary containing those same fields with the individual dictionaries in the list concatenated together and the adjacency_list indexes adjusted, as well as a tree_sizes list storing the size of each tree in the batch. Given a PyTorch Dataset object that returns tree data as a dictionary of tensors with the above keys, treelstm.batch_tree_input is suitable for use as a collate_fn argument to the PyTorch DataLoader object:

import treelstm

train_data_generator = DataLoader(
    TreeDataset(),
    collate_fn=treelstm.batch_tree_input,
    batch_size=64
)

Unbatching the batched tensors can be done via

torch.split(tensor, tree_sizes, dim=0)

Where tree_sizes is a list containing the number of nodes in each tree in the batch. This function is also provided by the treelstm.unbatch_tree_tensor function for convenience. As mentioned above, a tree_sizes list suitable for use by this function is generated by batch_tensors.batch_tree_tensor.

Example

Example code that generates tensors for the four node example tree above and trains a toy classification problem against the Tree labels is available in the example_usage.py script.