Руководство. Простая сверточная нейронная сеть (ConvNet)

Здесь мы создадим простую сверточную нейронную сеть (ConvNet) для классификации набора данных MNIST. Эта модель имеет простую архитектуру с тремя слоями обнаружения признаков (Conv -> ReLU -> MaxPool), за которыми следует финальный плотный слой, классифицирующий рукописные цифры MNIST. Обратите внимание, что эта модель, несмотря на свою простоту, должна достигать точности около 99 % после обучения в течение примерно 20 эпох. В этом примере сохраненная модель записывается в файл mnist_conv.bson. Кроме того, в нем демонстрируются основные принципы построения, обучения, сохранения модели, условного досрочного выхода и планирования скорости обучения. Для выполнения этого примера нам понадобятся следующие пакеты:

using Flux, MLDatasets, Statistics
using Flux: onehotbatch, onecold, logitcrossentropy, params
using MLDatasets: MNIST
using Base.Iterators: partition
using Printf, BSON
using CUDA
CUDA.allowscalar(false)

We set default values for learning rate, batch size, number of epochs, and path for saving the file mnist_conv.bson:

Base.@kwdef mutable struct TrainArgs
   lr::Float64 = 3e-3
   epochs::Int = 20
   batch_size = 128
   savepath::String = "./"
end

Данные

Чтобы обучить нашу модель, нужно собрать изображения вместе с их надписями и сгруппировать их в мини-пакеты (это ускоряет процесс обучения). Мы определяем функцию make_minibatch, которая принимает в качестве входных данных изображения (X) и их надписи (Y), а также индексы мини-партий (idx):

function make_minibatch(X, Y, idxs)
   X_batch = Array{Float32}(undef, size(X)[1:end-1]..., 1, length(idxs))
   for i in 1:length(idxs)
       X_batch[:, :, :, i] = Float32.(X[:,:,idxs[i]])
   end
   Y_batch = onehotbatch(Y[idxs], 0:9)
   return (X_batch, Y_batch)
end

make_minibatch выполняет следующие действия.

Создает массив X_batch размером 28x28x1x128 для хранения мини-пакетов.
Сохраняет мини-пакеты в X_batch.
Использует кодирование с одним активным состоянием для надписей изображений.
Сохраняет надписи в Y_batch.

get_processed_data loads the train and test data from Flux.Data.MNIST. First, it loads the images and labels of the train data set, and creates an array that contains the indices of the train images that correspond to each mini-batch (of size args.batch_size). Then, it calls the make_minibatch function to create all of the train mini-batches. Finally, it loads the test images and creates one mini-batch that contains them all.

function get_processed_data(args)
   # Load labels and images
   train_imgs, train_labels = MNIST.traindata()
   mb_idxs = partition(1:length(train_labels), args.batch_size)
   train_set = [make_minibatch(train_imgs, train_labels, i) for i in mb_idxs]

   # Prepare test set as one giant minibatch:
   test_imgs, test_labels = MNIST.testdata()
   test_set = make_minibatch(test_imgs, test_labels, 1:length(test_labels))

   return train_set, test_set

end

Модель

Now, we define the build_model function that creates a ConvNet model which is composed of three convolution layers (feature detection) and one classification layer. The input layer size is 28x28. The images are grayscale, which means there is only one channel (compared to 3 for RGB) in every data point. Combined together, the convolutional layer structure would look like Conv(kernel, input_channels => output_channels, ...). Each convolution layer reduces the size of the image by applying the Rectified Linear unit (ReLU) and MaxPool operations. On the other hand, the classification layer outputs a vector of 10 dimensions (a dense layer), that is, the number of classes that the model will be able to predict.

function build_model(args; imgsize = (28,28,1), nclasses = 10)
   cnn_output_size = Int.(floor.([imgsize[1]/8,imgsize[2]/8,32]))

   return Chain(
   # First convolution, operating upon a 28x28 image
   Conv((3, 3), imgsize[3]=>16, pad=(1,1), relu),
   MaxPool((2,2)),

   # Second convolution, operating upon a 14x14 image
   Conv((3, 3), 16=>32, pad=(1,1), relu),
   MaxPool((2,2)),

   # Third convolution, operating upon a 7x7 image
   Conv((3, 3), 32=>32, pad=(1,1), relu),
   MaxPool((2,2)),

   # Reshape 3d array into a 2d one using `Flux.flatten`, at this point it should be (3, 3, 32, N)
   flatten,
   Dense(prod(cnn_output_size), 10))
end

Чтобы соединить слои модели в цепочку, мы используем функцию Flux Chain. Она позволяет последовательно вызывать слои для заданных входных данных. Мы также используем функцию flatten для изменения формы выходного изображения из последнего слоя свертки. Наконец, мы вызываем функцию Dense, чтобы создать слой классификации.

Обучение

Прежде чем приступить к обучению модели, необходимо определить несколько функций, которые будут полезны в этом процессе.

augment добавляет случайный гауссовский шум к изображению, чтобы сделать его более устойчивым:
anynan проверяет, является ли какой-либо элемент параметра NaN или нет:
accuracy вычисляет долю входных данных x, правильно классифицированных сетью ConvNet:

augment(x) = x .+ gpu(0.1f0*randn(eltype(x), size(x)))
anynan(x) = any(y -> any(isnan, y), x)
accuracy(x, y, model) = mean(onecold(cpu(model(x))) .== onecold(cpu(y)))

Наконец, определим функцию train:

function train(; kws...)
   args = TrainArgs(; kws...)

   @info("Loading data set")
   train_set, test_set = get_processed_data(args)

   # Определим нашу модель.  Будем использовать простую сверточную архитектуру
   # с тремя итерациями слоев (Conv -> ReLU -> MaxPool), за которыми следует финальный плотный слой.
   @info("Building model...")
   model = build_model(args)

   # Загрузим модель и наборы данных в GPU, если эта возможность включена.
   train_set = gpu.(train_set)
   test_set = gpu.(test_set)
   model = gpu(model)

   # Перед запуском цикла обучения следует убедиться, что модель хорошо скомпилирована.
   model(train_set[1][1])

   # `loss()` вычисляет потери перекрестной энтропии между предсказанием `y_hat`
   # (вычисленным из `model(x)`) и эталонным истинным значением `y`.  Немного дополним данные,
   # добавив случайный гауссовский шум к изображению, чтобы сделать его более устойчивым.
   function loss(x, y)
       x̂ = augment(x)
       ŷ = model(x̂)
       return logitcrossentropy(ŷ, y)
   end

   # Обучим модель на основе заданного обучающего набора с помощью оптимизатора Adam
   # и по мере выполнения будем выводить результаты работы с тестовым набором.
   opt = Adam(args.lr)

   @info("Beginning training loop...")
   best_acc = 0.0
   last_improvement = 0
   for epoch_idx in 1:args.epochs
       # Обучаем в течение одной эпохи
       Flux.train!(loss, params(model), train_set, opt)

       # Завершаем при NaN
       if anynan(Flux.params(model))
           @error "NaN params"
           break
       end

       # Вычисляем точность:
       acc = accuracy(test_set..., model)

       @info(@sprintf("[%d]: Test accuracy: %.4f", epoch_idx, acc))
       # Если точность достаточно высока, выходим.
       if acc >= 0.999
           @info(" -> Early-exiting: We reached our target accuracy of 99.9%")
           break
       end

       # Если это лучшая точность, которую мы видели до сих пор, сохраняем модель.
       if acc >= best_acc
           @info(" -> New best accuracy! Saving model out to mnist_conv.bson")
           BSON.@save joinpath(args.savepath, "mnist_conv.bson") params=cpu.(params(model)) epoch_idx acc
           best_acc = acc
           last_improvement = epoch_idx
       end

       # Если мы не увидели улучшения за 5 эпох, снижаем скорость обучения:
       if epoch_idx - last_improvement >= 5 && opt.eta > 1e-6
           opt.eta /= 10.0
           @warn(" -> Haven't improved in a while, dropping learning rate to $(opt.eta)!")

           # Снизив скорость обучения, добавляем несколько эпох для улучшения результата.
           last_improvement = epoch_idx
       end

       if epoch_idx - last_improvement >= 10
           @warn(" -> We're calling this converged.")
           break
       end
   end
end

train calls the functions we defined above and trains our model. It stops when the model achieves 99% accuracy (early-exiting) or after performing 20 steps. More specifically, it performs the following steps:

Loads the MNIST dataset.
Builds our ConvNet model (as described above).
Loads the train and test data sets as well as our model onto a GPU (if available).
Defines a loss function that calculates the crossentropy between our prediction and the ground truth.
Sets the Adam optimiser to train the model with learning rate args.lr.
Runs the training loop. For each step (or epoch), it executes the following:
- Calls Flux.train! function to execute one training step.
- If any of the parameters of our model is NaN, then the training process is terminated.
- Calculates the model accuracy.
- If the model accuracy is >= 0.999, then early-exiting is executed.
- If the actual accuracy is the best so far, then the model is saved to mnist_conv.bson. Also, the new best accuracy and the current epoch is saved.
- If there has not been any improvement for the last 5 epochs, then the learning rate is dropped and the process waits a little longer for the accuracy to improve.
- If the last improvement was more than 10 epochs ago, then the process is terminated.

Тестирование

Наконец, чтобы протестировать модель, определяем функцию test:

function test(; kws...)
   args = TrainArgs(; kws...)

   # Загрузка тестовых данных
   _,test_set = get_processed_data(args)

   # Повторное построение модели с произвольными начальными весами
   model = build_model(args)

   # Загрузка сохраненных параметров
   BSON.@load joinpath(args.savepath, "mnist_conv.bson") params

   # Загрузка параметров в модель
   Flux.loadparams!(model, params)

   test_set = gpu.(test_set)
   model = gpu(model)
   @show accuracy(test_set...,model)
end

test загружает набор тестовых данных MNIST, повторно строит модель и загружает в нее сохраненные параметры (в mnist_conv.bson). Наконец, вычисляются предсказания модели для тестового набора и отображается точность теста (около 99 %). Полную версию этого примера можно посмотреть здесь: Simple ConvNets - model-zoo.

Ресурсы

Впервые опубликовано на сайте fluxml.ai 7 февраля 2021 г. Авторы: Эллиот Саба (Elliot Saba), Адарш Кумар (Adarsh Kumar), Майк Джей Иннес (Mike J Innes), Дхайрия Ганди (Dhairya Gandhi), Судханшу Агравал (Sudhanshu Agrawal), Самбит Кумар Даш (Sambit Kumar Dash), fps.io, Карло Лючибелло (Carlo Lucibello), Эндрю Динхобль (Andrew Dinhobl), Лилиана Бадильо (Liliana Badillo)