Efficient method for running Fully Convolutional Networks (FCNs)
By: Liad Pollak Zuckerman
Tags: Machine Learning, Deep Learning, Efficiency, Convolutional Network, Fully Convolutional Net

Fully convolutional (deep neural) networks, or FCNs, are commonly used for computer vision tasks, such as semantic segmentation, super-resolution, etc. One of their best properties is that they are applicable to inputs of any size, for example images of different sizes. However, running these networks on a large-scale input, such as high-resolution images or a video, can be exhaustive in terms of GPU memory consumption. In this post I will present a simple method to mitigate this problem. The proposed algorithm reduces GPU memory usage to only 3–30%.

What is a fully convolutional network?

A fully convolutional network (from now on: FCN) is a network consisted of only convolutional layers. If you’re unfamiliar with convolutional layers, here is a beginner’s guide. For simplicity, in this post we will focus on images, but the same applies to videos or other types of data.

A fully convolutional network for segmentation. input and output images from Nvidia dev blog

For many tasks in image processing, it is natural to require that the input and output images are of the same size. This could be achieved by using FCNs with proper padding (‘same’ in TensorFlow). Since this is a standard procedure, from now on we will assume it holds.

With this architecture, each pixel from the output image is a result of a calculation over a corresponding patch in the input image.

The size of this patch is called the receptive field (RF) of the network.

This is a key point. We’ll soon see how the algorithm uses this property

FCN with one layer: 3x3 convolution. each output pixel corresponds to a 3x3 patch in the input image

FCN with 2 layers: each is a 3x3 convolution. each output pixel corresponds to a 5x5 patch in the input image

Above we can see illustrations of FCNs with one (top) and two (bottom) 3x3 conv layers. In case of one layer (top), the blue output pixel on the right is a result of a calculation over the blue input patch on the left. When there are two layers (bottom) we have a feature map between the input and output. Each green pixel in the feature map (middle) is a result of calculating over a 3x3 green input patch (left), as in the one layer case. Similarly, each blue output pixel (right) is a result of calculating over a blue 3x3 feature map patch (middle), which originates in a 5x5 blue input patch (left).

So.. what’s the problem?

As mentioned before, theoretically, we could apply the network to any input size. Practically, however, computing the forward pass through the network requires keeping in memory huge feature maps, which exhaust GPU resources.

I encountered this limitation during my research on videos and 3-D images. The network simply would not run on our fancy NVIDIA v100 GPUs, which led me to develop this solution.

Efficient way to pass large inputs through an FCN

Traditional CNNs end with fully connected layers. Therefore, each output pixel is a result of a calculation over the entire input. This is not the case in FCNs. As we saw, only a receptive field sized patch from the input affects a single output pixel. As a result, to calculate a single output pixel there is no need to keep the entire feature map in memory!

In other words:

We can calculate small portions of the output values at a time, while passing only the necessary pixels from the input. This dramatically reduces the GPU’s memory usage!

Let’s have a look at the following example case:

Our input is a 28x28 image, as illustrated below. Due to memory limitations we can pass through our network patches up to size 12x12. The naive approach would be to input 12x12 patches from the original input (outlines in the figure below), and use the 12x12 output patches to construct the 28x28 output image.