Unlocking High-Resolution Computer Vision with Wafer-Scale Technology

[Updated April 2023: R1.8 of the Cerebras Software Platform now supports image segmentation on 50 megapixel images, up from 25 megapixels in R1.7.]

Convolutional Neural Networks have gained widespread popularity for various Computer Vision tasks in last decade. Accelerating CNNs that work on large images or volumes has been quite a challenging task for computer architects and system designers given the steep compute and memory requirements. At Cerebras, we’re uniquely positioned to accelerate large image CNNs, due to the architectural advantages of wafer-scale integration and our Weight Streaming execution mode. In this blog, we demonstrate our capabilities in this space and elaborate on these advantages.

Introduction

Breakthrough performance achieved with convolutional neural network models (CNNs) for computer vision is what triggered the current “AI summer”. The triggers were larger networks (more layers), new network designs (like ResNet), and fast hardware with which to train these.  CNN-based vision models have proved to be useful in autonomous vehicles and robots, agriculture, medicine, and other areas. Most recently, text-based image synthesis – also known as generative AI – based on CNN backbones (Dall-E, Stable Diffusion and Mid-Journey, for example) has been in the spotlight.

In many CV networks, there is a backbone model that generates feature maps. For example, ResNet serves as the backbone for the RetinaNet [1] object detection network, as shown in Figure 1. Feature maps from the backbone network are then consumed by a subsequent set of layers (task heads) that perform specialized computation. The backbone model, which is often pre-trained on a generic dataset, forms the foundation for the rest of the network. Some common backbones used in CV applications are the ResNet [3] and VGG [4] family of networks.

For image segmentation, the U-Net [2] architecture is often used. The goal of image segmentation is to partition an image into multiple segments or groups of pixels that correspond to different objects or regions. There are different types of segmentation depending on the requirements of the downstream task like semantic segmentation and instance segmentation. A modified form of U-Net is also used in the de-noising step of diffusion models, which are popular for text to image generation (like Stable Diffusion [5] from Stability AI).

Segmentation and classification for large images or volumes stress today’s AI compute systems. They need enough memory to store the weights and activations. For example, the popular 3D U-Net model, when training on 512³ sized volumes requires 64 gigabytes of storage for intermediate activations from the forward pass per image.

To grow memory and compute capacity, implementers have used large clusters of GPUs. But it has proven difficult to split a training job across large clusters and achieve an efficient, stable, successful training run.

The Cerebras Architecture and its advantages

At Cerebras, we believe our platform, the CS-2 system, is built using the only architecture in production that can overcome these challenges:

1. The Wafer-Scale Engine (WSE), which powers the CS-2 system, is the world’s largest computer chip. On a single CS-2, we can train a large vision model on large images. This eliminates the need to orchestrate and manage a large GPU cluster. The 40GB of on-chip SRAM and the high-speed, low-latency fabric provide ample memory capacity and performance, and interconnect bandwidth for handling the needs of CNNs. Monolithic training eliminates the problem of sub-linear performance scaling due to inter-chip synchronization, a problem common in large GPU clusters.
2. Our Weight Streaming mode of execution stores weights in an external device and streams the weights to the CS-2 on demand. This mode disaggregates compute performance from memory requirements, allowing the compute and the memory components to scale independently to effectively meet an application’s needs. And when one CS-2 is not fast enough or large enough for a challenging use case, we have demonstrated linear performance scaling across multiple CS-2s for data parallel training.

The advantages of the CS-2 present a promising opportunity to accelerate and scale CV models, both model size and image size, with minimal effort. Here are some key advantages for CV models:

• We can fit all the compute and activations required for 2D and 3D segmentation/classification models on a single WSE chip
• The Cerebras Graph Compiler for weight streaming allows scaling from small to large images (2D)/volumes (3D) on the same network architecture, with no extra changes required for work distribution, orchestration, or model code.
• Having all the operator data on a single chip means that bandwidth-bound operations such as normalizers aren’t slowed by having to move activation tensors in and out of the compute hardware from external memory.

We’ll say more about these points later.

How Do CNNs Work?

Before we delve into the details of how we implement CNNs on the Cerebras architecture, it’s worth briefly going over how CNNs work and why they are a good candidate for acceleration.

Convolution operators do most of the work in a CNN. Hence, understanding and optimizing the mapping and performance of a convolution operation is crucial.

Let’s start by thinking about an image as a 2D array of pixels, where each pixel is a single number.   A convolution takes in the image and produces a new image of, usually, the same size.    In the output image, the pixel value at a point is a weighted average of the input pixels in a small, normally square neighborhood centered at this same point.  The weights used are the same at every point.  Thus for example, the output pixel could be the average of the five pixels in the input at the given point and the four neighbors to the north, south, east, and west, minus half the average of the four neighbors in the northwest, northeast, southwest, and southeast directions. This would be a convolution using a 3-by-3 filter with filter coefficients.

You may worry about the output pixels at the edge. What about the neighboring points in the input image that lie outside the image boundary?  Some “padding” is used to supply values for these out-of-bounds input pixels; zero-padding is often used, in which these points are ignored.

In real CV applications, an image has another axis, called the feature axis.   Each pixel is represented as a vector of values, not a single number.   Initially, often it is the 3-vector of red, green, and blue intensities.   (At later, downstream layers of a CNN the image dimension, that is the number of pixels on each side, usually shrinks while the feature dimension grows.)  A convolution on such an image now computes the feature vector at an output point as a linear function of the neighboring point’s feature vectors.   Recall that for any linear function y = f(x) of a vector x, producing a vector y, there is a matrix A such that  y = A x, the product of the matrix A and the vector x. If y has dimension m and x has dimension n, then A is an m-by-n matrix; m rows and n columns.  Thus, in a 2D convolution with a 3-by-3 filter as above, the filter is defined by nine matrices.

Now imagine that we are working not on a single image but batch of them.   As this image batch passes through the network, we apply the convolution operators to them all together, making the feature vectors at a single point into, essentially, and matrix with as many columns as there are images in the batch.

Convolution therefore becomes the following in a 7-dimensional loop nest: