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

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

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

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

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)))

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

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