COBRA: E NHANCING DNN L ATENCY P REDICTION WITH L ANGUAGE M ODELS TRAINED ON S OURCE C ODE
A BSTRACT
With the recent developments of Deep Learning, having an accurate and device specific latency prediction for Deep Neural Networks (DNNs) has become important for both the manual and automatic design of efficient DNNs.
Directly predicting the latency of DNNs from their source code yields significant practical benefits.
It opens a way towards profilers that can instantly feedback the latency of a given piece of deep learning code to the developer.
In this paper, we conduct a preliminary study for source code based latency prediction of DNNs.
We introduce Co de B ased R untime A pproximation (COBRA), that leverages a transformer encoder to learn representations of short code snippets.
These representations are then aggregated by a Graph Convolutional Network (GCN) that captures the algorithmic dependencies and that estimates the latency of the implemented DNN.
Our experiments with COBRA show promising results and indicate that latency prediction from code can be competitive with traditional latency prediction methods for DNNs.
I NTRODUCTION
Deep learning methods have had tremendous success for a variety of applications, e g., image and automatic speech recognition (ASR), natural language processing (NLP), data restoration or generation.
A major challenge is the development of efficient DNN architectures that can also be deployed on mobile devices or on embedded platforms with restricted memory or computational power (Cai et al., 2019; Wu et al., 2019).
When tailoring DNN architectures to a specific target device, deep learning engineers typically require knowledge of the DNN latency.
Simply deploying a DNN and measuring its latency requires a lot of manual effort.
In particular, the deployment often relies on many steps such as graph optimization, operator fusion, weight and activation quantization or DNN compilation.
Many methods for DNN latency prediction have been proposed.
Among them, graph convolutional network (GCN) based predictors like proposed for BRP-NAS (Dudziak et al., 2020) are today considered to be the state-of-the-art.
They use a GCN to predict the latency of a DNN from its network graph representation.
Graph vertices are individual network layers and edges correspond to the data flow between the layers.
Each vertex is characterized by a feature vector, the layer representation.
A major problem is how to design good layer representations, such that they carry enough information to predict the latency of real-world DNN architectures.
BRP-NAS for example uses one-hot encodings of the network layer type.
This very simple approach works well with the special set of DNN architectures provided in the NAS-Bench-201 dataset (Dong & Yang, 2020).
However, in general, the latency of single network layers, and hence of whole DNNs, does not only depend on the layer types but also on the parameters of the layers, such as the shape of the input and output tensors.
Hence, this layer representation does not generalize well to more realistic DNN architectures.
We propose a novel and flexible method to predict the latency of DNNs using layer representations extracted from source code.
We call this method Co de B ased R untime A pproximation (COBRA).
First, COBRA interprets the source code to extract the corresponding network graph and layer implementations.
The layer implementations consists of small code snippets with function calls to a deep learning framework.
Second, it uses a transformer encoder to embed each layer implementation into a representation that is well suited for latency estimation.
We propose a special training trick for the transformer encoder that enables us to learn source code embeddings that are especially convenient for DNN latency prediction.
Finally, these extracted layer representations are aggregated by a GCN that captures the data dependencies between the function calls and estimates the latency of the DNN.
To our knowledge, COBRA is the first method that directly uses code representations to infer the latency of DNNs.
In our experiments, we compare COBRA with BRP-NAS, a GCN based predictor that uses hand-crafted layer representations that consist of the layer type, all the layer parameters, as well as of the measured layer latency.
Interestingly, COBRA consistently outperforms GCN based predictors that use such hand-crafted representations.
This paper is structured as follows: Section 2 gives an overview of important related work.
The details of the COBRA model and first experiments are given in Sections 3 and 4.
Finally, Sectiongives an outlook of our future directions of work.
R ELATED WORK
Language models on code: The recent works on CodeBERT (Feng et al., 2020) and Codex (Chen et al., 2021) have shown that we can train transformers like BERT (Devlin et al., 2019) and GPT- (Brown et al., 2020) as language models (LM) for source code.
These models work well for tasks involving both programming and natural language, such as code or docstring generation.
Further,
Allamanis et al.
(2021) used self-supervised learning to detect and repair code, given very simple defects.
However, to our knowledge it has not been tried yet to predict intrinsic properties of source code, such as the corresponding latency.
Latency predictors.
Different methods have been developed to predict the latency of DNN models for inference.
The simplest one consists of using the number of FLOPs as a proxy for the latency estimation.
However, the latency of a DNN has been shown not to be strongly correlated with
FLOPs and therefore the predictions are often inaccurate (Ma et al., 2018; Tan et al., 2019; Wu et al., 2019).
As a consequence, two distinct types of predictors have been devised to improve the quality of latency inference.
The first approach is to sum up the individual latencies of all the different elements composing a given neural network.
Thus, many works (Cai et al., 2019; 2020; Wu et al., 2019) simply use lookup tables to obtain the latency of each layer in the neural network and then derive the end-toend latency by summing the layer latencies.
However, the efficiency of this layer-wise predictor is limited.
Deep learning frameworks often fuse multiple layers into a single operation to accelerate the inference on real hardware devices.
This fusion is not taken into consideration by simple layerlevel predictors.
Besides, building a finite lookup table restricts us to select only a subset of all possible layer configurations.
To overcome such limitations, nn-Meter (Zhang et al., 2021) uses a different granularity by considering kernels instead of layers.
The authors define a kernel as being an independent execution unit on a device, i.e., the fusion of different layers, performed by the framework.
They develop an automatic method to detect the kernels of a neural network on a specific device and then build latency predictors in a supervised way for each of these kernels.
However, building many kernel latency predictors is costly.
The other class of latency predictors are end-to-end predictors, which directly estimate the latency of the DNN.
In contrast to sum-based predictors, end-to-end predictors consider the connectivity of the network.
BRP-NAS (Dudziak et al., 2020) uses a graph convolutional network (GCN).
Therefore, it requires the adjacency matrix of the given neural network to integrate the connectivity directly into the predictor.
HELP (Lee et al., 2021) uses a GCN or a multi-layer perceptron (MLP) to predict the latency of DNNs.
They test their approach on cell-based DNNs, like in the NAS-Bench- (Dong & Yang, 2020) dataset, and also on network architectures based on the FBNet (Wu et al., 2019) and on the MobileNetV3 (Howard et al., 2019) backbones.
HELP and BRP-NAS yield highly [Figure] accurate latency prediction for DNNs.
However, the DNN architectures that they consider for the training and testing are all very similar, in the sense that layers of the same type always share the same set of parameters and therefore have equal latency.
Of course, this is not true for general DNN architectures.
M ETHOD
As shown in Figure 1 COBRA performs three steps: 1) It interprets the code and extracts a set of layer implementations, as well as their call graph.
A transformer encoder that consumes layer implementations one at a time and embeds them into feature vectors.
A GCN that uses both the call sequence and the layer embeddings to compute an estimate of the latency of the defined DNN.
The interpreter is a rule-based module.
It identifies and extracts all layer implementations, that are single function calls to a deep learning toolbox.
It provides two outputs: 1) The set of layer implementations C
= { C 1 , C 2 , ..., C N } .
To simplify the problem we apply normalization, meaning that we parse all the function calls to the deep learning toolbox and re-order their parameters , such that they are always passed by name and in exactly the same order.
Further, all white spaces are removed.
An adjacency matrix A that encodes the dependencies between all C n .
More specifically, for a DNN that is defined by N layer implementations, A is an N × N upper triangular matrix.
The transformer encoder processes each C n individually.
First, C n is tokenized.
We use the uncased English BERT tokenizer from the PyTorch transformer implementation that uses a dictionary of D = 30 , 522 tokens.
The transformer encoder itself has an architecture similar to BERT (Devlin et al., 2019) and also uses the same positional encoding.
More specifically, it consists of 4 multihead attention blocks with 16 attention heads each.
Both the hidden and embedding dimension of the attention blocks are C = 512 .
Layer normalization is always performed before attention is applied.
The maximum sequence length is K = 600 tokens.
Consequently, the output of the encoder for the input C n is H n ∈ R K × C .
Figure 2 shows our pre-training setup for the transformer encoder.
It uses a combination of selfsupervised and supervised training to learn embeddings that are suited for latency prediction.
For the self-supervised part the encoder is pre-trained to predict masked input tokens.
We denote these predictions with ˆ E .
This way, the transformer encoder learns to capture the properties of dependence in single layer implementations and learns to act as a language model for source code.
Input tokens are predicted using a single linear layer of size C × D that is applied to the output H n of the encoder.
We noticed that it is problematic that layer implementations rely on many parameters like input and output shapes or kernel sizes.
Although these parameters heavily influence the layer latencies, the transformer will not be able to capture the properties of dependence between them if we do not provide additional information.
In particular, these parameters are all independent and it is impossible to predict them once they are masked.
To solve this problem, we append the tokenized measured latency of the layer implementations.
With this additional information, all parameters of a [Figure] layer that have an influence on the latency become dependent.
These properties of dependence can then be learned by the encoder.
Note that the tokenized latency is only appended during pre-training of the transformer, but not for fine tuning or for inference.
For the supervised part, the transformer encoder is also trained to explicitly predict the layer latency.
Therefore, H n is processed with a latency prediction head.
Its input is the average of the token embeddings x n = H Tn 1 /K ∈ R C .
More specifically, use a MLP with one hidden layer of sizeand an output layer of size 1 that estimates the log-latency of the layer.
Consequently, the transformer encoder is trained with two losses, i.e., the categorical cross-entropy loss for token reconstruction and the mean squared error (MSE) for the log-latency estimation.
By combining these two losses for pre-training, we make sure that the transformer encoder learns to extract token representations that capture the syntax of the layer implementation while also carrying information about their latency.
The GCN combines the output of the transformer encoder for all C n ∈ C and calculates a latency estimate for the full DNN.
In general, a GCN (Kipf & Welling, 2017) can process graph signals.
Let
G =
( V, E ) be a graph, where V is a set of N vertices and E is a set of edges connecting them.
A graph signal assigns a vector x i
= x ( v i ) ∈ R C to each vertex.
GCNs exploit the properties of dependence between all x i in order to make predictions on the vertex or on the graph level, while taking the graph structure into account.
For a graph with N vertices, the graph signal can be summarized in a matrix X =
[ x 1 , x 2 , ..., x N ] T ∈ R N × C .
The layer-wise propagation rule of the GCN is then defined as
[Equation]
where l = 1 , 2 , ..., L , X ( l ) and W ( l ) are the layer index, the layer activation and the weight matrix, respectively.
Further, A ∈ R N × N is the adjacency matrix of the graph.
In our case, v i is a DNN layer implementation consisting of function calls to a deep learning toolbox and x i is the corresponding layer representation that is extracted by the transformer encoder.
In particular, the input of the GCN is X (0) =
X .
The adjacency matrix is a binary upper triangular matrix with A i,j = 1 if layer i is connected to layer j and A i,j = 0 otherwise.
σ is the ReLU activation function.
The latency estimates are computed from X ( L ) the activation of the last GCN layer.
More specifically, we first perform a global average pooling over the vertices dimension X ( L ) T 1 /N in order to get a vector that encodes the properties of the whole graph and then apply a linear layer to obtain the output.
In our experiments, the GCN is trained in a supervised setup using the mean absolute percentage error (MAPE) E [ | ˆ y n − y n | /y n ] as a loss.
Moreover, we train it independently of the transformer encoder, meaning that the parameters of the encoder are kept fixed after pre-training and are not optimized together with the GCN parameters.
E XPERIMENTS
We perform our experiments using Sony’s deep learning framework NNabla (Hayakawa et al., 2021).
It provides a static computational graph that we parse in order to extract the needed layer implementations and the adjacency matrix.
[Table] [Figure] plot of BRP-NAS++, because it is almost similar to BRP-NAS.
Especially for low-latency networks the performance of COBRA is very good.
Although the test set also contains networks with very large latency, we only show low latency models here, because the prediction errors are most pronounced for them.
Datasets:
We collect two automatically generated datasets, one to pre-train the transformer encoder and the other to train the COBRA model.
The first dataset consists of 291 , 658 randomly generated layer implementation C n and their measured latency t n .
The set is split into 289 , 198 training and 2 , 460 test samples.
We consider the layers Add , Convolution , Linear , ReLU , BatchNormalization , AveragePooling ,
MaxPooling and GlobalAveragePooling , where all the input arguments are chosen randomly, allowing various layer parameterizations.
For Linear , we can for example choose random input and output tensor shapes.
Note, that the dataset contains each layer parameterization only once, meaning that all layers contained in the test set have never been seen before.
This way, we can use it to test the inter- and extrapolation properties of the transformer encoder for latency prediction.
More details are given in Section A.1.1.
The latency for each random parameter configuration is measured on an NVIDIA 1080TI GPU with 11GB VRAM, using NNabla graph profiling functionality.
The second dataset is used to train the GCN and is constructed from randomly generated DNN code and the DNN latencies that have also been measured on an NVIDIA 1080TI GPU.
In total, we collect a dataset of 2000 random ResNet-like models (He et al., 2016) and split the dataset into 1350 training, 150 validation and 500 test samples.
ResNet-like means that we not only allow skipconnections between neighbouring, but also between any two arbitrary network layers.
The number of network layers is random, but is never larger than 226 .
Also the individual layer types and their configuration, e.g., the number of output channels and the stride of a convolution, are chosen randomly.
For more details about the degrees of freedom that we allow, please refer to Section A.2.
Figure 4 shows the histogram of the latencies in the collected datasets.
Most of the collected layer implementations have a low latency.
More specifically, 20% of them have a latency that is lower than [Figure] [Figure] a y e r i m p l e m e n t a ti on s
[Figure]
1 ms and 80% have a latency that is lower than 4 .
3 ms.
However, the latency distribution is heavily tailed.
In particular, the maximum latency of a single layer is 1524 .
2 ms.
Further, most collected
DNNs yield a latency inferior to 200 ms.
However, similar to single layer implementations, the corresponding latency distribution is heavily tailed.
Performance of COBRA:
In our first experiment, we compare the performance of COBRA to three other methods for DNN latency prediction, namely: 1)
The layer-wise sum that estimates the latency of a DNN by summing up the measured latencies of the individual network layers.
It is also calibrated by a scaling factor to fit the latencies in the training set.
Although this method has access to the measured latencies of all network layers of a DNN, it is known to perform poorly.
Hence, it acts as a lower baseline.
A GCN based predictor that uses hand-crafted layer representations, consisting of the one-hot encoded layer type and all the layer parameters.
This representation is proposed by Zhang et al.
The GCN architecture is similar to the state-of-the-art Dudziak et al.
(2020), hence we call this method BRP-NAS. 3)
The GCN based predictor which uses layer representations that also contain the measured latency of each individual layer.
We call this method BRP-NAS++.
Although latency measurements of single layers can be costly to obtain in practice, we add this feature to see how good a GCN predictor can perform if it is provided with all available information.
To train the predictors, we first search for optimized hyperparameters on the validation set, using
Optuna (Akiba et al., 2019).
The best hyperparameters are summarized in Table 5 in the appendix.
Each predictor is then trained 10 times, starting from different random initializations and using both the training and the validation set.
Table 1 summarizes the root mean square error (RMSE), the mean absolute percentage error (MAPE), as well as the error bound at 1% , 5% , 10% and 25% of each latency prediction method.
More specifically, the error bound at k % measures for how many samples from the test set the latency is predicted with less than k % error (higher is better).
COBRA clearly outperforms layerwise sum, BRP-NAS and BRP-NAS++, when comparing the
MAPE and the error bounds.
In particular, COBRA yields less than 5% error for 96 . 2% of the test networks.
Compared to BRP-NAS++ that yields less than 5% error for only 92 .
7% of the test networks, COBRA is better by 3 . 5% .
For less than 1% error, COBRA outperforms BRP-NAS++ by 8 . 1% .
Note that the layerwise sum yields the lowest RMSE of all the methods.
That is because BRPNAS, BRP-NAS++ and COBRA are trained to minimize the MAPE, which measures the average relative prediction error.
In practice, estimators that achieve a small relative prediction error are more desirable than those that achieve a low absolute prediction error.
In particular, we want highly accurate estimates for low latency networks and allow larger estimation errors for networks with [Table] BatchNormalization [Table]
AveragePooling
[Table]
GlobalAveragePooling [Table] All [Table]
with
[Table]
without latency
Average
[Table]
a large latency.
This behaviour is not captured by the RMSE, because it measures the absolute prediction errors.
Performance of the transformer encoder:
In our second experiment, we evaluate how accurate the transformer encoder alone can predict the latency of single layers from their source code implementation.
Table 2 gives the latency prediction on the test set.
We observe that the transformer encoder can accurately predict the latency of unseen layer implementations, and can interpolate and extrapolate between layers of very different parameterizations.
Note, that it yields the lowest prediction errors for simple layers like ReLU , Add or Linear , where an error of less than 25% is obtained for 100% of the test samples.
Further, the transformer encoder struggles to predict the latency of convolutional layers.
A convolution is defined by the input tensor shape, the number of channels, the kernel size, the stride, the padding, the group, and the dilation.
Compared to all the other considered layers, it has the largest number of parameters and therefore is the hardest to predict.
In our last experiment, we show why it is important to append the layer latency to the layer implementation, when pre-training the transformer encoder.
More specifically, we train the transformer encoder twice, once with and once without the appended layer latency.
For both, we report the average perplexity (PPL) for masked token reconstruction, assuming that only parameter tokens are randomly masked.
This directly reflects how good the transformer encoder can capture the properties of dependence between the parameters, conditioned on a certain latency.
As shown in Figure 5, the perplexity is lower if the layer latency is appended to the layer implementation.
As argued before, all layer parameters are independent if we provide no latency information.
Hence, it is impossible to predict them from the masked layer implementation.
However, if we apply our training trick and also provide latency information, all layer parameters that influence the latency are conditionally dependent.
This enables us to learn properties of dependence between the layer parameters and the layer latency.
In particular, we can learn how masked parameters can be computed from others.
Hence, it also yields a much lower perplexity.
D ISCUSSION AND C ONCLUSIONS
Our experiments demonstrated the feasibility of latency prediction directly from source code.
In particular, we demonstrated that using a simple training trick, we can train a standard transformer encoder to predict DNN layer latencies from their source code implementation.
Further, we could successfully use the layer embeddings generated by the transformer encoder to improve state-of-theart GCN based latency predictors for DNNs.
Our method COBRA has two advantages: 1) Compared to BRP-NAS, it uses learned, rather than hand-crafted features to encode individual function calls.
Therefore, it is easy to extend it to new layer types.
Compared to BRP-NAS++, it does not require the measured latency of individual function calls, while still yielding a comparable performance.
Limitations of the current experiments: We are aware of some limitations of our current experimental setup that must be addressed in the future.
Our current training and validation data is automatically generated and therefore heavily normalized.
In particular, it does not capture the variability of a human authored source code, where there is for example lots of variations to the choice of variable names and how parameters are passed.
Nevertheless, our experiments are very encouraging and demonstrate the feasibility of latency prediction from code.
A study with human written code has to be conducted in the future to test the real-world performance of the system.
Due to time constraints, we only compared our method to BRP-NAS so far.
Although BRP-NAS has proven to be a very good latency predictor for DNNs, fair comparisons to other state-of-the-art latency predictors like nn-Meter must be performed.
Up to now, we also consider latency predictions for inference on a GPU.
Therefore, the experiments must be extended to other platforms, such as embedded devices or mobile DNN accelerators.
Future directions of research: During our work on COBRA, we identified two research areas for machine learning on code that are in our opinion of general interest and therefore should be addressed in the future:
1) Compared to natural language, the behaviour of code heavily depends on numeric values such as bounds for counters in loops, indices when accessing data or parameters that are used for basic calculations.
We think that current tokenizers from NLP typically lack the ability to represent these values in a suited way, what makes them not particularly well suited for machine learning on code.
This issue was also raised by Geva et al.
A simple way to address this issue could be to represent numbers with positional encodings.
COBRA can most likely be improved, by using such ideas.
Another important challenge is dealing with extremely long sequences.
As already discussed in the introduction, information about the behaviour of code is very distributed and can span multiple files and hundreds or even thousands of lines of code.
Transformers cannot easily be scaled up to process such long sequences because the computational complexity of the attention mechanism grows quadratic with the sequence length.
Therefore, we need good ways to scale these models up.
One approach that we also used in this work is to split the code into fragments.
Each fragment would be processed individually and the results are combined.
However, we can also think of some kind of hierarchical processing with structured attention, where the attention is first computed across single lines of code, then across code blocks and finally across files.
A A PPENDIX A. T RAINING DETAILS A.1. T RANSFORMER ENCODER
The transformer encoder is defined in a similar way to the BERT model (Devlin et al., 2019).
The tokenizer and the positional encoding are the same, but the architecture is a lighter version of BERT BASE , in particular with fewer layers and a smaller hidden size.
All the hyperparameters of the transformer are summarized in Table
[Table] [Table]
We want to make the transformer learn the non-linearity dependencies between a specific layer implementation and its corresponding latency.
We therefore generate a dataset composed of code snippets, where each snippet is the code implementation required to execute a layer with a specific configuration.
To keep it simple
[Table]
we use random sampling to select the layer configurations, i
[Table]
e., for each sample we pick uniformly a random layer configuration on a specific range of values.
The layers dataset is therefore very general, with various possible configurations and allowing even to have configuration with non-equal input width and height.
Table 4 describes the different layer configurations as well as the number of samples for each type of layer.
For each layer configuration, we also measure its latency.
A.1. GCN
We observe that a proper training of the different GCNs heavily depends on the values of the hyperparameters.
We therefore use a hyperparameter optimization framework called Optuna (Akiba et al., 2019) to find good hyperparameters for each GCN.
The list of all the hyperparameters is available in Table 5.
In particular we observe that a batch size equal to 2 is giving the best performance on the test loss.
A. DNN DATASET
We explain here more in details how the dataset of DNN architectures is built.
The connections between layers are encoded by the adjacency matrix A .
We consider therefore the set of all the binary upper triangular matrices of fixed size N × N .
Then, any operation of the set of predefined operations defined in Table 6 can be associated to a specific layer, with the objective of allowing various configurations of layers.
The first layer and last layer are consecrated respectively to the input and the output.
One way to sample from this architecture space is just to draw uniformly an architecture from it, i.e., one adjacency matrix A and one operation per layer.
However, to control the density of the drawn
[Table]
GlobalAveragePooling
[Table] [Table]
architectures, we introduce the probability p of having two arbitrary layers connected.
Therefore large values of p involve more dense networks whereas p = 0 .
5 is uniform sampling.
We fix p = 0 .to
avoid having almost only very dense networks.
Once we have drawn the adjacency matrix and the operations of each layer, we perform a pruning step similar to in NAS-Bench-101 (Ying et al., 2019) to ensure that the resulting DNN is valid.
First, we discard adjacency matrices which correspond to invalid architectures because of the input node and the output node being not connected.
In addition,
[Table]
we remove useless dangling nodes, leading to architectures with very different sizes and therefore a wider range of latency than other dataset such as NAS-BENCH 201 (Dong & Yang, 2020).
To avoid sampling the same architecture multiple times, we apply a graph hashing algorithm detecting whether two architectures are isomorphic or not (Ying, 2019).
Finally, since we can add multiple layers together with different input shapes, we add convolution blocks (Conv+BatchNorm+ReLU) before adding layers to decrease the resulting input shape of each layer to the greatest common divisor among the different input shapes.