An Efficient End-to-End 3D Voxel Reconstruction based on Neural Architecture Search

With the development of 3D deep learning, there is a growing body of work studying implicit neural representation and reconstruction. Here we investigate two sub-directions closely related to our work. The first sub-direction in implicit neural representation and reconstruction is the prediction of SDFs from dense samples using MLP networks [14] [15] [16] [4] [5] [17] [18]. For example, DeepSDF [14] early proposes to learn and reconstruct continuous SDFs for a category of 3D objects using an MLP network. FFN [15] maps Fourier features and learns high-frequency functions in low-dimensional domains to improve the fitting capability of MLP. To overcome the difficulty of fitting SDFs of one category of objects, NI [4] firstly proposes to overfit the global shape of an individual object with a small MLP network. This method actually implements a lossy compression of a 3D model by storing MLP parameters. However, the default network with 7553 parameters may fail in reconstruction, especially for complex objects. NGLOD [5] learns the latent vectors of octree vertices in different levels of details (LOD) to predict local SDFs of an object. Although the reconstruction accuracy can be improved obviously, the storage capacity in this local-fitting method is greatly increased as the latent vectors of a large number of LOD vertices needs to be stored.

The second sub-direction in implicit neural reconstruction is the prediction of occupancy of voxels using encoder-decoder networks [3] [19] [20]. Similar to ONet [3] reviewed in Sec. I, IM-Net [19] learns the generative models of shapes for one category of objects. CONet [20] combines convolutional encoders with implicit occupancy decoders to represent detailed reconstruction of objects and 3D scenes. However, the reconstruction accuracy of these works is obviously lower than the SDF prediction networks. Moreover, the network size of these works is much larger than most of MLP networks.

In theory, predicting occupancy of points is much easier than predicting their SDF values with the same network and training data. Therefore, in this paper, we adopt this idea to directly predict the occupancy of binary voxels.

Neural Architecture Search (NAS) methods essentially aim to provide an automated way to design architectures as an alternative to manual architectures. Our work is closely related to reinforcement learning based NAS work [12] [13] [21] [22] [23]. For example, owing to the weight sharing idea, ENAS [23] can significantly reduce the computational power required to traverse the search space [21]. Although later researchers propose a different framework DARTS  [24] [25] [26], this kind of approach is not suitable for MLP architectures search.

In addition, the process of NAS needs to accurately evaluate the performance of each network architecture. A straightforward solution is to train an architecture from scratch and then test it on the validation dataset, which is very time consuming. Instead of accurately evaluating the network architecture on the target task, the researchers proposed the proxy task method. The proxy task means training on subsets of dataset or fine-tuning with fewer epochs [23] [27] [28] [29]. Although these methods improve the speed of NAS, a rough evaluation inevitably treats some promising network architectures as poor networks. A post-processing step after NAS is proposed in this paper. As a result, those potential networks with fewer network parameters can be found.

Binary voxels can be directly visualized as one of the explicit representations of 3D objects. In general, voxels inside an object are defined as 1, while voxels outside the object are defined as 0. Training a neural network to classify the binary voxels of a given model in 3D space enables end-to-end neural representation and reconstruction. Therefore, the post-processing steps, such as surface reconstruction using marching cubes [4] [5], can be avoided.

Objects are normalized in a 3D space denoted by $\mathcal{H}\!=\![-1,1]^{3}$. The entire normalized space is divided equally into $N^{3}$ parts, each of which is regarded as a voxel. The set of voxels inside the object is denoted by $\mathcal{V}$. For each voxel $\mathbf{p}$ in $\mathcal{H}$, the neural network $f_{\theta}(\cdot)$ outputs the probability of $\mathbf{p}\in\mathcal{V}$, which is between 0 and 1. To optimize the parameters $\theta$ of the neural network, the cross-entropy classification loss $\mathcal{L}(\theta)$ is used as the following,

Most of previous works employ MLP as their global or local SDF fitting network. Therefore, the number of layers, the number of neurons in each layer, and the activation functions in each layer are chosen as our neural architecture search space.

Specifically, in order to reduce the search time and take into account the reasonable search range, the search space for the number of nodes is {8, 12, 16, 20, 24, 28, 32, 40, 48, 56, 64}. Moreover, we let the controller decide the activation function for each layer, making the network architecture more expressive. The search space of activation functions is {$ReLU,ELU,Swish$} [30] [31] [11], each of which can be represented by

In the above selection of activation functions, the traditional activation functions $Sigmoid$ and $Tanh$ are not included. This is because we experimentally find that adding these activation functions will reduce the classification performance of the network. As stated in the previous NAS works [32] [33] [21], the choice and design of the search space plays a crucial role in the success of NAS. Ablation experiments shown in Sec. IV-D demonstrate the performance of adding activation functions into the search space.

In order to search a ‘proper’ neural architecture specifically for a given object, we utilize a mature NAS algorithm ENAS [23]. The searching process is drawn in Fig. 2. The controller samples the MLP layers from the search space consisting of different numbers of nodes and activation functions. After an MLP network is trained, its reward can be obtained on the validation set. The reward is then fed back to the controller RNN for policy-gradient descent.

Unlike ENAS, the search strategy in our method focuses on how to choose a better MLP network with different numbers of neurons and activation functions. However, the raw reward in ENAS only considers the classification score, and the internal controller always generates network architectures that make the classification reward higher. Since one of the expectations of the neural reconstruction is to minimize the number of network parameters [4], the impact of network size should be factored into the reward. Thus, the reward is designed as

It is worth noting that if the controller generates an output layer in the process of generating MLP layers, the number of MLP layers will no longer increase. This explains that the number of MLP layers can also be searched during the NAS process.

After completing the NAS process, we introduce a post-processing step to select smaller neural architectures. There are two reasons for this step. First, to speed up the NAS process, we use a proxy task [21] that is similar to NAS-FPN [34] to shorten the training time of the target task. We train the proxy task $e_{1}$ epochs instead of $e_{2}$ epochs used to train the target network. This early termination method speeds up the convergence time of rewards by a factor of $e_{1}/e_{2}$. Since the proxy task does not train the network to converge, the ranking of the network accuracy cannot represent the final ranking of the network. Therefore, the network with accuracy slightly lower than the highest accuracy can be considered as candidates.

Second, although we have used the network size reward term in Eq. 3 to influence the network chosen by NAS, a fixed weight is not applicable to different objects. A network with a little lower network size and classification accuracy may be excluded, compared to the network with the least reward in the NAS process. Post-processing selection of multiple candidates would greatly alleviate this problem.

Specifically, we filter out all candidate networks whose accuracy is lower than the highest accuracy in the NAS process up to a threshold $t$ ($t\!=\!0.1\%$ in experiments). Then the network with the smallest size will be selected. The ablation experiments shown in Sec. IV-C validate the effectiveness of the proposed post-processing step.

In data processing, 3D models are firstly normalized in a sphere with the radius $0.9$. Then we utilize PyMesh library to voxelize a model to $N^{3}$ resolution (N=128 by default).

In voxel sampling, we first sample all surface voxels and their outer layer voxels as the support samples of classification boundary. Then, we down-sample $1/4$ other non-support voxels and copy the support voxels to the same number. Thus, the total sample number $K$ is $1/2$ the number of other voxels, which is roughly equivalent to 1M.

In the NAS configuration, the maximum number of total network layers and hidden layers is 8 and 6, respectively. The controller samples 6 MLP layers of different architectures at a time. Due to the weight sharing strategy proposed in ENAS, we terminate the search by only sampling 5 times.

Datasets we used in this paper include Thingi10K [9], Thingi32 and ShapeNet150. Thingi10K is composed of $10,000$ 3D-printing models, which have been tested in NI [4]. NGLOD [5] mainly verifies two other small datasets: Thingi32 and ShapeNet150. Thingi32 contains 32 simple shapes in Thingi10K. ShapeNet150 contains 150 shapes in the ShapeNet dataset [35], including 50 cars, 50 airplanes, and 50 chairs.

The metrics for evaluation are common 3D intersection over union (IoU) and Chamfer distance (CD). The former metric 3D IoU is defined as the ratio of the intersection and the union voxels of a reconstructed model and its ground truth model. The latter metric CD is defined as the bi-directional minimum distance [36] from the surface voxels $\mathcal{S}_{r}$ of a reconstructed model to the surface voxels $\mathcal{S}_{r}$ of its ground truth model, which is expressed by

We compare our approach with one voxel reconstruction method ONet [3], and two SDF prediction methods NI [4] (with the default configuration) and NGLOD [5] (with LOD level 1). First, the comparison is conducted on Thingi32 and ShapeNet150 datasets. The experimental results of the previous three methods are shown in the upper part of Table I, and our results are shown in the last row. Note that the results of NI-Thingi32 and ONet-ShapeNet150 are obtained through the official trained networks. The results of NI-ShapeNet150 and NGLOD on two datasets are obtained by our training using the official codes. Since ONet is trained on one category of objects, it cannot be used for objects in Thingi32.

Since Thingi32 only contains models of simple shapes, NI, NGLOD and our method perform similarly, and they can all reconstruct these models without noticeable errors. For ShapeNet150, the performance differences of the four methods are easily distinguishable. Since there are some complicated objects in ShapeNet150, ONet and NI cannot reconstruct them correctly. This results in a large increase in the metric CD. NGLOD can handle more objects correctly than ONet and NI, and gain better CD and IoU. The proposed method shows significant improvement in CD and IoU, which indicates that our method does not suffer as much performance degradation as other methods when dealing with complex objects. Fig. 3 depicts three reconstructed voxels (one for each category in ShapeNet150) for all four methods as a visual comparison.

As ShapeNet150 contains only 150 objects in three categories, we further conduct another experiment on Thingi10K dataset. We train the NAS networks for $1,000$ models randomly selected from Thingi10K, and compare the reconstruction results with NI. Fig. 4 depicts the histograms of CD and IoU for the two methods, respectively. The proposed method is obviously superior to NI.

This ablation experiment aims to observe the influence of removing NAS, the proposed size reward, and the post-processing step. The experimental results are shown in the lower part of Table I (rows $5$ to $7$). The $5$-th row shows our results without NAS. The network here is much the same as NI. The main difference is that instead of predicting SDF values, we directly classify voxels, and over-sample the support voxels. As NI successfully reconstructs the 3D models in Thingi32, the method in the $5$-th row achieves only a slight improvement. However, for the more complicated ShapeNet150 dataset, the improvements of IoU and CD are significant. This means that binary classification is much easier to be fitted by a same network than SDF prediction, which experimentally validates the first theoretical problem described in Sec. I.

The $6$-th row shows our results using NAS, but without adding network size to reward and the post-processing step. Since there is no scheme to control the network scale, the NAS obtains networks with more parameters at average. Owing to larger network parameters, the method in the $6$-th row gains the best IoU and CD performance on two datasets.

The $7$-th row shows our results using NAS and size reward, but without the post-processing step. Since the reward is improved to control network size, the mean number of network parameters is roughly the same to it in the $5$-th row, and is reduced by $\sim$$15\%$ compared to the $6$-th row. As a result, the performance of the method in the $7$-th row degrades slightly.

The advantage of the post-processing step is reflected in the last row. Compared to the $7$-th row, although the performance is promoted slightly, the mean size of the networks is significantly reduced. This verifies the effectiveness of the proposed method in selecting a suitable network architecture.

In order to choose suitable activation functions as the candidates in our search space, we test five activation functions, which are Sigmoid, Tanh, ReLU, ELU, and Swish, on Thingi32 and ShapeNet150 separately. This is also an ablation experiment, as NAS only searches different numbers of layers and nodes at the moment. There is no change to the setting except that only one activation function is used during NAS. The experimental results are shown in Table II.

The activation function ReLU, ELU or Swish, can achieve good accuracy, while the results of Sigmoid are dramatically worse, especially on ShapeNet150. Although the mean network size of Sigmoid appears to be minimal, we experimentally find that the reconstruction for Sigmoid is likely to fail. In order to reduce invalid searches, we remove Sigmoid and Tanh from the final search space of activation functions. Compared to using one fixed activation function, using three activation functions {ReLU, ELU, Swish} as the search space reduces CD and IoU very slightly, but searches obviously smaller networks.