Tree DNN: A Deep Container Network

MTL provides the facility to solve multiple tasks using one DNN model. Few MTL studies share parameters among their tasks and are known as shared trunk methodology as proposed by Lu et al., (2022) Xie et al., (2022) Ma et al., (2018). All the MTL methods work in a similar manner, where input sample is inferenced through multiple operations and operations are shared among tasks as proposed by Ott et al., (2022). This way, it reduces efforts in training a task from scratch by utilizing the trunk as shown by Santiago et al., (2021). There are multiple applications like Keceli et al., (2022) proposed a MTL-based framework for identifying each plant’s type and disease by considering raw image and transformed feature-maps from another pretrained DNN.

However, MTL methods do not provide solution, if each task needs training from a different dataset.

TreeDNN is a container network that is a collection of multiple DNN models prepared to serve the tasks independently. As the name suggests a portion of the DNN model is shared by all the tasks and is referred as trunk of the tree as shown in Figure 1. Each branch of the tree is dedicated to a particular task. In this work, TreeDNN utilizes the ops of vision model. Trunk of the tree is designed to receive generalized training through all the datasets. TreeDNN is designed to speed up the switching among the tasks, which are different branches of the tree therefore, branches of the tree need to be lighter (on scale of parameters). The speedup is attained because of keeping trunk of the tree on the working memory and consuming the time only in loading the particular branch of the tree. Whole process can be segregated among a) Model designing b) Dataset preparation and c) Training.

(a) Model Designing: A typical vision model is a combination of operations like Convolution, BatchNorm, Relu, Pooling, Linear, Fully connected layers etc. A TreeDNN architecture contains a trunk, branches and both are designed with the above mentioned vision operations. In this work, we have considered MobileNet-V2 as a backbone architecture therefore trunk of the Tree is mainly a collection of inverted residual blocks. Branch of the TreeDNN is prepared from the second half portion of Mobilenet-V2 and the number of layers are reduced/ increased based on the complexity of the task performed at that particular branch. The process of model creation is defined through line 25-30, in Algorithm 1. It is interesting to realize that $M_{i}:Trunk\rightarrow{Branch_{i}}$ is a fully-functioning DNN model capable to serve $Task_{i}$.

(b) Dataset Preparation: TreeDNN requires joint training through multiple datasets. In order to maintain trade-off between bias and variance while performing training at each branch, it is required to design the minibatch considering all the datasets. Since minibatch is prepared with the participation from each of the datasets, we refer to it as $Fed\_Batch$ (federated batch). The preparation of the federated batch is shown in line:3, Algorithm 1.

(c) Training TreeDNN is trained in two steps: i) Generalized Training ii) Specialized Training.

• •

  Generalized Training: In first step, all the branches and trunk of the tree are trained together. During training, the samples are inferenced through the trunk and branches of the tree. This way, loss is calculated at each branch and then all the losses are aggregated as shown in equation 1. Cumulative loss is then backpropagated to the network through branch and the trunk equation 2. This step helps in training the trunk of the tree. Where $L$ is the loss value, $\hat{y}_{j}$ is the observed output at each branch and $y_{j}$ is the label at branch $j$.

  $$J(w^{T},b)=\sum_{i=1}^{n}L_{Task:1}(\hat{y}_{1}^{i},y_{1}^{i})+\sum_{i=1}^{n}L_{Task:2}(\hat{y}_{2}^{i},y_{2}^{i})+\ldots+\sum_{i=1}^{n}L_{Task:k}(\hat{y}_{k}^{i},y_{k}^{i})$$

  (1)

  $$J(w^{T},b).backward()$$

  (2)

• •

  Specialized Training: Now, it is required to tune each branch corresponding to a specific task. This step is performed by backpropagating the loss back to the network while keeping the gradient of trunk False equation 3. Specialized training is applied at each branch independently as shown in equation 4. Where $I$ is the input samples.

  $$\hat{v}=Trunk(I)\;,\;\;\;\;\;gradient(Trunk)=False\\$$

  (3)

  $$\hat{y}_{j}=branch_{j}(\hat{v})\;,\;\;\;\;\;Jj(w^{T},b)=\sum_{i=1}^{n}L_{Task:j}(\hat{y}_{j}^{i},y_{j}^{i})$$

  (4)

TreeDNN is developed with the agenda of solving the issues raised by loading and inferencing through multiple camera models or related systems. TreeDNN is a type of multi-task learning which supports the training through multiple datasets in a joint fashion. TreeDNN provides the flexibility of keeping the trunk of the tree on the working memory and because of that only a branch of the tree is required to be invoked at the time of execution. Hence, it reduces the loading time of the model, which can be witnessed in Table 1. The trunk of the tree works as the initial layers for any of the tasks and is shared among all the tasks therefore it helps in reducing the storage requirement of the initial layers’ parameters. Table 1 showcases the case study of camera models and the performance of the TreeDNN when utilized in a real environment having 8 DNNs. The specified tasks for these models could be De-mosaic, Auto-Focus, Auto-Exposure, Auto-White balance, Auto-ISO, Color correction, Brightness enhancement etc.

The performance of the TreeDNN in the present experiment can be seen in Table 2. We have considered CIFAR10, CIFAR100, Caltech101, Caltech256, SVHN datasets to train the TreeDNN with 5 branches. SOTA results are the performance of MobileNet-V2 on aforementioned datasets when trained separately. It can be realized from the results that difficult datasets like Caltech256 and Caltech101 has shown improvement in performance because of utilizing the generalized trunk in the tree which works as regularization while training. However, there is a minute drop in the performance of relatively easy datasets. TreeDNN can be deployed as a single model for all the tasks which also opens the direction of mounting new pretrained branch in an already grown TreeDNN to increase the scalability of the model for new tasks. Where the decision about selection of most appropriate branch can be taken by considering the task similarity.