A mixed-scale dense convolutional neural network for image analysis

Problem Definition.

In this paper, we apply our approach to real-valued 2D images. We define an image as a set of pixels 𝒙∈ℝm×n×c with m rows, n columns, and c channels. We denote the image corresponding to a single channel j of x as 𝒙j. Many image-processing problems can be written as the problem of finding a function f that takes a certain image x and produces an output image y; i.e., f:ℝm×n×c→ℝm′×n′×c′. Note that the dimensions of the output image can be different from those of the input image. In image classification problems, for example, the output image consists of a single probability value for each of the c′ possible classifications; i.e., m′=n′=1. In the rest of this paper, however, we focus on problems with dense outputs, i.e., with the number of rows and columns of the output image identical to those of the input image: m′=m and n′=n, similar to “pixel to pixel” architectures (16).

Convolutional Neural Networks.

Convolutional neural networks (CNNs) model the unknown function f by using several layers that are connected to each other in succession. Each layer i produces an output image 𝒛i∈ℝmi×ni×ci, called a feature map, using output of the previous layer 𝒛i−1 as input. The dimensions of the layer output 𝒛i can be different from those of the layer input 𝒛i−1. The input image x is taken as the first layer 𝒛0, and the final layer produces the output image y.

Each individual layer can consist of multiple operations. A common layer architecture first convolves each channel of the input feature map with a different filter, then sums the resulting convolved images pixel by pixel, adds a constant value (the bias) to the resulting image, and finally applies a nonlinear operation to each pixel. These operations can be repeated using different filters and biases to produce multiple channels for the output feature map. Thus, the output 𝒛ij of a single channel j of such a convolutional layer is given by𝒛ij=σ(gij(𝒛i−1)+bij).[1]Here, σ:ℝmi×ni→ℝmi×ni is a nonlinear operation such as the popular sigmoid function or rectified linear unit (ReLU) (7), bij∈ℝ is the bias, and gij:ℝmi−1×ni−1×ci−1→ℝmi×ni convolves each channel of the input feature map with a different filter and sums the resulting images pixel by pixel,gij(𝒛i−1)=∑k=0ci−1C𝒉ijk𝒛i−1k,[2]where C𝒈𝒂 is a 2D convolution of image a with filter g. Different ways of handling the boundaries of the image during convolution are possible: Here, we use reflective boundaries. Often, the filters 𝒉ijk are relatively small (e.g., 3×3 pixels), enabling faster computation of network outputs and making the network easier to train. The architecture of the final layer can differ from other layers and can depend on the application: Common choices include using a fully connected layer instead of a convolutional one (2) or using a softmax function as the nonlinear operation for classification problems (5). A schematic of a two-layer CNN architecture is shown in Fig. 1.

• Download figure
• Open in new tab
• Download powerpoint

Fig. 1.

A schematic representation of a two-layer CNN with input x, output y, and feature maps 𝒛1 and 𝒛2. Arrows represent convolutions with nonlinear activation.

The goal of training a CNN is to find filters 𝒉ijk, biases bij, and potential other parameters, such that the CNN performs the task that is required. In supervised learning, training is achieved by using a set of Nt representative inputs 𝑿={𝒙^1,…,𝒙^Nt} with corresponding correct outputs 𝐘={𝒚^1,…,𝒚^Nt} and iteratively minimizing a chosen error metric between Y and the CNN output for X. Because of the specific architecture of CNNs, partial gradients of the error with respect to the filters and biases can be computed accurately and efficiently through backpropagation for several popular error metrics, enabling the use of efficient gradient-based optimization algorithms (17).

DCNNs.

DCNNs use a network architecture similar to standard CNNs, but consist of a larger number of layers, which enables them to model more complicated functions. In addition, DCNNs often include downscaling and upscaling operations between layers, decreasing and increasing the dimensions of feature maps to capture features at different image scales. Many DCNNs incrementally downscale feature maps in the first half of the layers, called the encoder part of the network, and subsequently upscale in the second half, called the decoder part. Skip connections are often included between feature maps of the decoder and encoder at identical scales (5). A schematic representation of a common encoder–decoder DCNN architecture is shown in Fig. 2.

• Download figure
• Open in new tab
• Download powerpoint

Fig. 2.

A schematic representation of a common DCNN architecture with scaling operations. Downward arrows represent downscaling operations, upward arrows represent upscaling operations, and dashed arrows represent skip connections.

In general, the increased depth of DCNNs compared with shallow CNNs makes training more difficult. The increased depth often makes it more likely that training gets stuck in a local minimum of the error function and can result in gradients that become either too large or too small (13). Furthermore, DCNNs typically consist of many parameters (e.g., filters and biases), often several million or more, that have to be learned during training. The large parameter space can make training more difficult, by increasing both training time (18) and the likelihood of overfitting the network to the training data (8), thereby forcing large training sets. Several additions to standard DCNN architectures have been proposed, including batch normalization layers (13), which rescale feature maps between layers to improve the scaling of gradients during training; highway connections (14); residual connections (15); fractal networks (19), which allow information to flow more easily through deep networks by skipping layers; and dropout layers (8), in which feature maps are randomly removed from the network during training, reducing the problem of overfitting large networks.

Although these additions have advanced image processing in several fields (12), they can be difficult to routinely apply in areas such as biomedical imaging and materials science. Instead, traditional imaging algorithms are used, such as the Hough transform (20) and template matching (21), or manual processing [e.g., biological image segmentation (22)].

Mixing Scales.

Instead of using downscaling and upscaling operations to capture features at different scales, the MS-D architecture uses dilated convolutions. A dilated convolution D𝒉,s with dilation s∈ℤ+ uses a dilated filter h that is nonzero only at distances that are a multiple of s pixels from the center.^∗ Recently, it was shown that dilated convolutions are able to capture additional features in DCNNs that use the traditional scaling approach (23). Furthermore, instead of having each layer operate at a certain scale as in existing DCNNs, in the mixed-scale approach each individual channel of a feature map within a single layer operates at different scale. Specifically, we associate the convolution operations for each channel of the output image of a certain layer with a different dilation:gij(𝒛i−1)=∑k=0ci−1D𝒉ijk,sij𝒛i−1k.[3]The proposed mixed-scale approach alleviates many of the disadvantages of the standard downscaling and upscaling approach. First, large-scale information about the image quickly becomes available in early layers of the network through relatively large dilations, making it possible to use this information to improve the results of deeper layers. Furthermore, information at a certain scale can be used directly to inform decisions at other scales without having to pass through layers at intermediate scales. Similar advantages were recently found when training large multigrid architectures (24). No additional parameters have to be learned during training, since the mixed-scale approach does not include learned upscaling operations. This results in smaller networks that are easier to train. Finally, although dilations sij must be chosen in advance, the network can learn which combinations of dilations to use during training, making identical mixed-scale DCNNs applicable across different problems (experiments below).

Dense Connections.

When using convolutions with reflective boundaries, the mixed-scale approach has an additional advantage compared with standard scaling: All network feature maps have the same number of rows and columns as the input and output image, i.e., mi=m and ni=n for all layers i, and hence, when computing a feature map for a specific layer, we are not restricted to using only the output of the previous layer. Instead, all previously computed feature maps {𝒛0,…,𝒛i−1}, including the input image x, can be used to compute the layer output 𝒛i. Thus, we change the channel image computation 1 and the convolutional operation 3 to𝒛ij=σ(gij({𝒛0,…,𝒛i−1})+bij)gij({𝒛0,…,𝒛i−1})=∑l=0i−1∑k=0cl−1D𝒉ijkl,sij𝒛lk.[4]Similarly, to produce the final output image y, all feature maps can be used instead of only those of the last layer. We call this approach of using all previously computed feature maps densely connecting a network.

In a densely connected network, all feature maps are maximally (re)used: If a certain useful feature is detected in a feature map, it does not have to be replicated in other layers to be used deeper in the network, as in other DCNN architectures. As a result, significantly fewer feature maps and trainable parameters are required to achieve the same accuracy in densely connected networks compared with standard networks. The smaller number of maps and parameters makes it easier to train densely connected networks, reducing the risk of overfitting and enabling effective training with relatively small training sets. Recently, a similar dense-connection architecture was proposed which relied on a relatively small number of parameters (25); however, in ref. 25 the dense connections were used only within small sets of layers at a single scale, with traditional downscaling and upscaling operations to acquire information at different scales. Here, we combine dense connections with the mixed-scale approach, enabling dense connections between the feature maps of the entire network, resulting in more efficient use of all feature maps and an even larger reduction of the number of required parameters.

MS-D Neural Networks.

By combining mixed-scale dilated convolutions and dense connections, we can define a DCNN architecture that we call the MS-D network architecture. Similar to existing architectures, an MS-D network consists of several layers of feature maps. Each feature map is the result of applying the same set of operations given by Eq. 4 to all previous feature maps: dilated convolutions with 3×3 pixel filters and a channel-specific dilation, summing resulting images pixel by pixel, adding a constant bias to each pixel, and finally applying a ReLU activation function. The final network output is computed with the same set of operations applied to all feature maps, using 1×1 pixel filters instead of 3×3 pixel filters. In other words, channels of the final output image are computed by taking linear combinations of all channels of all feature maps and applying an application-specific activation function to the result:𝒚k=σ′(∑i,jwijk𝒛ij+bk′).[5]Different ways of choosing the number of channels per layer are possible. Here, we use a simple approach with each layer having the same number of channels, denoted by the network width w, and the number of noninput and nonoutput layers of the network denoted by the network depth d. A graphical representation of an MS-D network with w=2 and d=3 is shown in Fig. 3. The parameters that have to be learned during training are the convolution filters 𝒉ijkl and biases bij of Eq. 4 and the weights wijk and biases b′k of Eq. 5. Given a network depth d and width w and number of input channels c𝑖𝑛 and output channels c𝑜𝑢𝑡, the number of trainable parameters N𝑝𝑎𝑟=N𝑓𝑙𝑡𝑠+N𝑤𝑔𝑡𝑠+N𝑏𝑖𝑎𝑠 is given by N𝑓𝑙𝑡𝑠=9∑i=0d−1w(iw+c𝑖𝑛), N𝑤𝑔𝑡𝑠=(wd+c𝑖𝑛)c𝑜𝑢𝑡, and N𝑏𝑖𝑎𝑠=wd+c𝑜𝑢𝑡.

• Download figure
• Open in new tab
• Download powerpoint

Fig. 3.

Schematic representation of an MS-D network with w=2 and d=3. Colored lines represent 3×3 dilated convolutions, with each color representing a different dilation. Note that all feature maps are used for the final output computation.

Compared with existing DCNN architectures, the MS-D network architecture has several advantages. Due to the mixing of scales through dilated convolutions and dense connections, MS-D networks can produce accurate results with relatively few feature maps and trainable parameters. Furthermore, an MS-D network learns which combination of dilations to use during training, allowing the same network to be effectively applied to a wide variety of problems. Finally, all layers are connected to each other in the same way and computed using the same set of standard operations, making MS-D networks easier to implement, train, and use in practice. MS-D networks do not include learned scaling operations or advanced layer types to facilitate training and do not require architecture changes when being applied to different problems. These advantages can make MS-D networks applicable beyond semantic segmentation, with potential value in classification, detection, instance segmentation, and adversarial networks (16).

Setup.

We implemented the MS-D architecture in Python, using PyCUDA (26) to enable GPU acceleration of computationally expensive parts such as convolutional operations. We note that existing frameworks such as TensorFlow (9) or Caffe (10) typically do not support the proposed mixed-scale approach well, since they assume that all channels of a certain feature map are computed in the same way. Furthermore, existing frameworks are mostly optimized for processing large numbers of relatively small images by efficiently implementing convolutions using large matrix multiplications (27). To allow the application of MS-D networks to problems with large images, we implemented the architecture using direct convolutions. Computations were performed on two workstations, with an NVidia GeForce GTX 1080 GPU and four NVidia Tesla K80 GPUs, respectively, all running CUDA 8.0.

In general, deeper networks tend to produce more accurate results than shallower networks (2). Because of the dense connections in MS-D networks, it is possible to effectively use networks that have many layers and few channels per layer, resulting in very deep networks with relatively few channels. Such very deep networks might be more difficult to train than shallower networks, as explained above. However, we did not observe such problems and were able to use the extreme case of each layer consisting of only one channel (w=1) and the number of layers d controlling the number of trainable parameters. We initialize all convolution filter parameters based on the same considerations as in ref. 3 by sampling random values from a zero-mean normal distribution with a SD of 2/nc, where nc is the number of incoming and outgoing connections of a feature map channel: nc=9(c𝑖𝑛+w(d−1))+c𝑜𝑢𝑡. All other trainable parameters are initialized to zero. Finally, in most experiments we use equally distributed dilations sij∈[1,10] by setting the dilation of channel j of layer i equal to sij=((iw+j)mod10)+1.

In segmentation problems with L labels, we represent correct outputs by images with L channels, with channel j set to 1 for pixels that are assigned to label j and set to 0 for other pixels. We use the soft-max activation function in the final output layer and use the ADAM optimization method (17) during training to minimize the cross-entropy between correct outputs and network outputs (5). To compare results of MS-D networks with existing architectures for segmentation problems, we use the global accuracy metric (4), defined as the percentage of correctly labeled pixels in the network output, and the class accuracy metric (4), computed by taking the average of the true positive rates for each individual label.

Simulated Data.

In a first experiment, network input consist of 512×512-pixel single-channel images of objects with two shapes (circles and squares), 3 different sizes, and 6 possible textures, with added Gaussian noise. From all 36 combinations of shape, size, and texture, we train networks to detect 6 specific combinations, e.g., large squares with a horizontal texture, small circles with a diagonal texture, etc. We chose this segmentation problem because it requires DCNNs to combine features at small scales (pixel intensity and texture) with features at larger scales (size and shape) to produce accurate results. An example input is shown in Fig. 4A, with colors indicating the six combinations that have to be detected in Fig. 4B. We compare segmentation results of trained MS-D networks with those of the popular U-Net architecture (5): We use a TensorFlow implementation (28). U-Net architectures are similar to that shown in Fig. 2. Two main parameters influence performance: the number of downscaling (and subsequent upscaling) operations and the number of channels per feature map. We train each network with the same set of 105 randomly generated images, using a batch size of one image, and stop training once global accuracy for a different set of 100 validation images has not improved for 105 iterations.

• Download figure
• Open in new tab
• Download powerpoint

Fig. 4.

(A–C) Example of the segmentation problem of the simulated dataset, with (A) the single-channel input image, (B) the correct segmentation with labels indicated by color, and (C) the output of a trained MS-D network with 200 layers.

In Fig. 5, class accuracy for an independent test set of 100 images is shown as a function of the number of trainable parameters for MS-D networks with w=1 and d∈{25,50,100,200} layers and U-Net networks with two, three, four, and five scaling operations and various numbers of channels. The performance of the U-Net networks depends significantly on the chosen number of scaling operations. Networks with three scaling operations are able to achieve around 80% accuracy with relatively few parameters, but do not improve significantly when using more channels per feature map, while networks with four scaling operations are able to achieve around 95% accuracy, but require a large number of parameters to do so. For a given number of parameters, MS-D networks are able to achieve significantly higher accuracies than all tested U-Net architectures, especially with relatively few parameters, and the performance of MS-D networks is similar for different choices of dilations.

• Download figure
• Open in new tab
• Download powerpoint

Fig. 5.

The class accuracy of a set of 100 simulated images (Fig. 4) as a function of the number of trainable parameters for the proposed MS-D network architecture and the popular U-Net architecture. For each U-Net network (U-Net-q), q indicates the number of scaling operations used. For the MS-D architecture, results are shown for dilations sij∈[1,10] (solid line) and sij∈{1,2,4,8,16} (dashed line).

CamVid Dataset.

Next, we compare results for the CamVid dataset (29), using 367 training, 101 validation, and 233 testing color images of road scenes with 360×480 pixels (4). The goal is to segment 11 classes such as cars, roads, sidewalks, and pedestrians. We train MS-D networks and U-Net networks with local contrast-normalized images (30) until there is no improvement in global accuracy of the validation set, using minibatches of 10 images for MS-D networks and smaller minibatches of 3 images for U-Net networks due to memory constraints. We also report results for the SegNet architecture (4), showing the two best global accuracy results from table 1 of ref. 4, and two traditional segmentation methods (31, 32), showing the two best results from table 2 of ref. 4. For U-Net networks, the number of feature map channels was chosen such that the number of parameters was similar to that of the SegNet.

Table 1 shows global and class accuracies. MS-D segments with highest global and class accuracy, while using roughly 10 times fewer parameters. Furthermore, an MS-D network with 100 layers achieves similar accuracies to other network architectures while using 30–40 times fewer parameters.^† Fig. 6 shows global accuracy during training for validation and training sets, for both the U-Net network and an MS-D network. Lack of improvement for the U-Net network in validation set accuracy and its difference with training set accuracy indicate overfitting of the chosen training set. Due to the smaller number of trainable parameters, the MS-D network improves validation set accuracy for more training iterations, with a significantly smaller difference with training set accuracy, showing reduced risk of overfitting of MS-D networks and the ability to accurately train with relatively small training sets. In addition, MS-D networks are able to achieve accurate results without pretraining additional large datasets, e.g., ImageNet (11), or relying on large pretrained networks, e.g., VGG (2).

• Download figure
• Open in new tab
• Download powerpoint

Fig. 6.

The global accuracy of a U-Net network and an MS-D network as a function of the training epoch for the CamVid dataset. Given are the accuracies for the validation set (solid lines) and the training set (dashed lines).

Segmenting Biomedical Images.

To test whether an MS-D network can be easily applied to a new problem without adjustments, we use the same network parameters as above, with w=1, d=100, and dilations sij∈[1,10] applied to segmenting cell structures. We use eight manual segmentations of 512×512×512 tomographic reconstructions of (mouse) lympoblastoid cells, consisting of five labels: nuclear envelope, euchromatin, heterochromatin, mitochondria, and lipid drops. A sample tomographic slice and corresponding manual segmentation are shown in Fig. 7 A and B. The labeling of cell structures depends on multiple factors at different image scales, such as the position of the structure relative to other structures, and the pixel intensity differences between two structures can be relatively small, making it difficult to use traditional methods to perform automatic labeling. Instead, researchers rely on time-consuming manual segmentation.

• Download figure
• Open in new tab
• Download powerpoint

Fig. 7.

(A–C) A tomographic slice of the test cell (A), with the corresponding manual segmentation (B) and output of an MS-D network with 100 layers (C).

To learn limited 3D features, we use five channels in the input image of the MS-D network: the current slice to be segmented and four adjacent slices. Of eight manual cell segmentations, we randomly chose six for training and one for validation and report results for the remaining cell. During training, we used a batch size of 10 images and stopped after no improvements in global accuracy for the validation cell, yielding network parameters with the best global accuracy. Fig. 7C shows network output for the slice in Fig. 7A, showing high similarity to manual segmentation. The remaining differences between network output and manual segmentation, indicated by an arrow in Fig. 7, typically represent ambiguous cell structure (see Figs. S1 and S2 for additional results). Final global accuracy and class accuracy of the trained network for the test cell are 94.1% and 93.1%, indicating that identical MS-D networks can be trained for different problems. Results for two other challenging problems are given in Figs. S3 and S4.

Denoising Large Tomographic Images.

Finally, we use the above architecture, changing only the nonlinear function of the final layer from the soft-max function to the identity, and train on the different task of denoising tomographic reconstructions of a fiber-reinforced minicomposite. A total of 2,160 images of 2,5602 pixels were reconstructed using 1,024 acquired X-ray projections to obtain images with relatively low amounts of noise (Fig. 8A). Noisy images of the same object were obtained by reconstructing using 128 projections (Fig. 8B). The input is a noisy image, with a corresponding noiseless image used as target output during training. From the sample top, 500 images were used for training, and 100 images were used for validation. Fig. 8C shows output for a tested image near the sample bottom, computed in 2.05 s using a GTX 1080 GPU (see Fig. S5 for additional timings). The MS-D network accurately denoises highly noisy images by learning image features from the training set, and identical MS-D networks can be easily applied to different problems with minimal changes.

• Download figure
• Open in new tab
• Download powerpoint

Fig. 8.

(A–C) Tomographic images of a fiber-reinforced minicomposite, reconstructed using 1,024 projections (A) and 128 projections (B). In C, the output of an MS-D network with image B as input is shown. Bottom Right Insets in A–C show enlarged images of small regions indicated by red squares.