This example demonstrates how to train the VGG network against the Imagewoof image classification dataset (a harder version of Imagenette).
A VGG-16 network is instantiated from the ImageClassificationModels library of standard models, and applied to an instance of the Imagewoof dataset. A custom training loop is defined, and the training and test losses and accuracies for each epoch are shown during training.
As a note: the current implementation of the Imagewoof dataset loads all images into memory as floats, which can lead to memory exhaustion on machines with less than 16 GB of available RAM.
VGG, as a network, is extremely sensitive to training tweaks and batch size. This code sample demonstrates using a custom decreasing learning rate schedule, similar to that employed in the original paper.
To begin, you'll need the latest version of Swift for
TensorFlow
installed. Make sure you've added the correct version of swift to your path.
To train the model, run:
cd swift-models
swift run -c release VGG-Imagewoof