Build a CNN for Image Classification with Swift for TensorFlow

Swift for TensorFlow (S4TF) was a high-level API developed by Google that let you use the TensorFlow machine learning framework from the Swift programming language, pairing TensorFlow's features with Swift's type safety and expressiveness. It was an ambitious project — first-class automatic differentiation in the language — but Google archived it in 2021, so the code here no longer runs on a maintained toolchain.

If you landed here to actually build a CNN today, jump to the modern alternatives below. The step-by-step tutorial that follows documents how S4TF expressed a CNN, for readers studying its design.

In this tutorial, we walk through building a Convolutional Neural Network (CNN) for image classification using the Swift API for TensorFlow: importing modules, loading and preprocessing image data, defining the CNN architecture, creating the model, training it, and evaluating its performance.

1. Importing necessary modules and libraries:

   • Import the TensorFlow module to access the Swift API for TensorFlow functionalities:

     import TensorFlow

2. Loading and preprocessing image data:

   • Load the image data and preprocess it using techniques like resizing, normalizing pixel values, and converting images to tensors. For example:

     let dataset = Dataset(...)
     let resizedImages = dataset.images.resized(to: (224, 224))
     let normalizedImages = resizedImages / 255.0
     let tensorImages = Tensor<Float>(normalizedImages)
     let labels = Tensor<Int32>(dataset.labels)
     let trainDataset = Dataset(elements: (tensorImages, labels))

3. Defining the CNN architecture:

   • Define the structure of the CNN model using the building blocks provided by the Swift API for TensorFlow. Configure convolutional layers, pooling layers, and fully connected layers. For example:

     struct CNN: Layer {
         var conv1 = Conv2D<Float>(filterShape: (3, 3, 3, 32), padding: .same, activation: relu)
         var conv2 = Conv2D<Float>(filterShape: (3, 3, 32, 64), padding: .same, activation: relu)
         var flatten = Flatten<Float>()
         var dense = Dense<Float>(inputSize: 7 * 7 * 64, outputSize: 10, activation: softmax)
      
         @differentiable
         func callAsFunction(_ input: Tensor<Float>) -> Tensor<Float> {
             let convolved1 = conv1(input)
             let convolved2 = conv2(convolved1)
             let flattened = flatten(convolved2)
             return dense(flattened)
         }
     }

4. Creating the model:

   • Create an instance of the CNN model using the defined architecture. Initialize the model, specify the optimizer, and choose an appropriate loss function for image classification. For example:

     var model = CNN()
     let optimizer = Adam(for: model)
     let loss = softmaxCrossEntropy(logits:reduction:)
      

5. Training the model:

   • Train the CNN model using labeled training data. Perform forward and backward propagation, apply gradient descent optimization, and update the model's parameters. For example:

     let epochs = 10
     for epoch in 1...epochs {
         var epochLoss: Float = 0
         for batch in trainDataset.batched(batchSize) {
             let (images, labels) = (batch.first, batch.second)
             let gradients = gradient(at: model) { model -> Tensor<Float> in
                 let logits = model(images)
                 let batchLoss = loss(labels: labels, logits: logits)
                 epochLoss += batchLoss.scalarized()
                 return batchLoss
             }
             optimizer.update(&model.allDifferentiableVariables, along: gradients)
         }
         print("Epoch \(epoch): Loss: \(epochLoss)")
     }

6. Evaluating the model:

   • Evaluate the performance of the trained model using a separate set of labeled test data. Calculate metrics like accuracy, precision, and recall to assess how well the model generalizes to new, unseen images.

   For example:

   let testImages = loadTestImages()
   let testLabels = loadTestLabels()
   let testTensorImages = Tensor<Float>(testImages)
   let testTensorLabels = Tensor<Int32>(testLabels)
   let testDataset = Dataset(elements: (testTensorImages, testTensorLabels))
    
   var correctPredictions = 0
   var totalPredictions = 0
   for batch in testDataset.batched(batchSize) {
       let (images, labels) = (batch.first, batch.second)
       let logits = model(images)
       let predictions = logits.argmax(squeezingAxis: 1)
       correctPredictions += predictions .== labels
       totalPredictions += predictions.shape[0]
   }
   let accuracy = Float(correctPredictions) / Float(totalPredictions)
   print("Test Accuracy: \(accuracy)")

The code above shows S4TF's design: a model is a struct conforming to Layer, and @differentiable plus the gradient(at:) function gave automatic differentiation as a language feature. That design is interesting history, but it will not compile on a maintained toolchain.

What to use instead

Since S4TF is archived, build image classifiers with a maintained framework:

• Python + TensorFlow/Keras — the direct successor path. The same CNN in a few lines:

  import tensorflow as tf
   
  model = tf.keras.Sequential([
      tf.keras.layers.Conv2D(32, 3, activation="relu", padding="same",
                             input_shape=(224, 224, 3)),
      tf.keras.layers.MaxPooling2D(),
      tf.keras.layers.Conv2D(64, 3, activation="relu", padding="same"),
      tf.keras.layers.MaxPooling2D(),
      tf.keras.layers.Flatten(),
      tf.keras.layers.Dense(10, activation="softmax"),
  ])
  model.compile(optimizer="adam",
                loss="sparse_categorical_crossentropy", metrics=["accuracy"])

• Python + PyTorch — the other mainstream choice, subclassing nn.Module.

• On Apple platforms — use Create ML (train with a GUI/Swift API) or export a Keras/PyTorch model to Core ML with coremltools and run it on-device. This is the supported way to do ML in Swift today.

Related guides

• Getting started with PyTorch
• AI/ML interview & engineering hub

Sources

• Swift for TensorFlow — archived repository (read-only)
• Core ML — Apple Developer