Deep neural networks use deeper and broader structures to achieve better performance and consequently, use increasingly more GPU memory as well. However, limited GPU memory restricts many potential designs of neural networks. In this paper, we propose a reinforcement learning based variable swapping and recomputation algorithm to reduce the memory cost, without sacrificing the accuracy of models. Variable swapping can transfer variables between CPU and GPU memory to reduce variables stored in GPU memory. Recomputation can trade time for space by removing some feature maps during forward propagation. Forward functions are executed once again to get the feature maps before reuse. However, how to automatically decide which variables to be swapped or recomputed remains a challenging problem. To address this issue, we propose to use a deep Q-network(DQN) to make plans. By combining variable swapping and recomputation, our results outperform several well-known benchmarks. Limited GPU memory restricts model performance due to two different reasons. Firstly, there is a trend that deep neural networks (DNNs) use deeper and more GPU memory-intensive structures (Wang et al., 2018) , and have continuously made improvement in various computer vision areas such as image classification, object detection, and semantic segmentation (He et al., 2016a; Simonyan & Zisserman, 2014; Krizhevsky et al., 2012; Ronneberger et al., 2015; Goodfellow et al., 2016; Szegedy et al., 2015) . Likewise, empirical results show that deeper networks can achieve higher accuracy (He et al., 2016b; Urban et al., 2016) . Deeper network means higher consumption of GPU memory. Secondly, He et al. (2019) shows that bigger input batch size can speed up the training process and achieve higher accuracy. However, a bigger input batch size requires more GPU memory to store intermediate variables. We want more GPU memory to get better performance. The rationale to utilize CPU memory by offloading, and later prefetching variables from it is twofold. Firstly, the size of the CPU memory is usually bigger than that of GPU memory. If we do not use variable swapping, all the tensors will stay in GPU memory. Figure 1 shows the details of variable swapping. Secondly, due to the availability of the GPU direct memory access (DMA) engines, which can overlap data transfers with kernel execution. More specifically, a GPU engine is an independent unit which can operate or be scheduled in parallel with other engines. DMA engines control data transfers, and kernel engines can execute different layer functions of DNNs. Hence, in the ideal case, we can completely overlap DNNs training with variable swapping. Therefore, variable swapping is efficient. Regarding recomputation, some feature maps are not stored in GPU memory in forward propagation, but the feature maps are gotten by running forward functions in backpropagation, as shown in Figure 2 . Why do we combine swapping with recomputation? Because recomputation uses GPU computing engines to reduce memory usage, and variable swapping uses DMA engines to save memory. Different engines can run parallelly. If we execute recomputation during data transfers, we will not waste computing engines or DMA engines. It is hard to decide which variables should be swapped or recomputed. Different DNNs have different structures. Networks have thousands of variables during training, so it is intractable to enumerate the search space exhaustively. Some existing works use heuristic search algorithms for recompu- Figure 1 : The upper graph shows GPU operations in a standard neural network in a time sequence. The lower one shows how to add variable swapping operations. The nodes in the same column represent they occur at the same time. We copy X 0 into CPU memory while reading X 0 . After data transfer and reading, we f ree X 0 from GPU memory. Before using X 0 again, we allocate space for X 0 and transfer it back to GPU memory. Figure 2: If we do not store X 1 in the memory in the forward propagation, we need to execute the layer 0 forward function again to get X 1 for the layer 1 backward function. tation or swapping with limited information from computational graphs. For example, they do not consider the time cost of recomputing different layers or swapping different variables. Additionally, they do not make plans for recomputation during swapping in order to increase GPU utilization. Our work utilizes more information from computational graphs than theirs and makes plans automatically for users. The contribution of our paper is that we propose a DQN algorithm to make plans for swapping and recomputation to reduce memory usage of DNNs. Users only need to set memory usage limits and do not require background knowledge on DNNs. Additionally, the variable swapping and recomputation will not decrease the accuracy of networks. In this paper, we propose a DQN to devise plans for variable swapping and recomputation to reduce memory usage. Our work can work well with different memory limits. Our method provides plans automatically for users. They only need to set a memory limit and do not require background knowledge on DNN or machine learning algorithm. Our method can work well for different network structures such as ResNet, VGG, K-means, SD ResNet, and LSTM. Besides, the variable swapping and recomputation do not decrease the accuracy of networks.
We propose a reinforcement learning based variable swapping and recomputation algorithm to reduce the memory cost.
{"src": "scitldr-AIC_train", "paper_id": "BJxg7eHYvB"}