Глубокая сверточная генеративно-состязательная сеть (DCGAN)

using Pkg

# Активируем новую среду проекта в текущем каталоге
Pkg.activate(".")
# Добавляем в среду необходимые пакеты
Pkg.add(["Images", "Flux", "MLDatasets", "CUDA", "Parameters"])

using Base.Iterators: partition
using Printf
using Statistics
using Random
using Images
using Flux: params, DataLoader
using Flux.Optimise: update!
using Flux.Losses: logitbinarycrossentropy
using MLDatasets: MNIST
using CUDA

Base.@kwdef struct HyperParams
    batch_size::Int = 128
    latent_dim::Int = 100
    epochs::Int = 25
    verbose_freq::Int = 1000
    output_dim::Int = 5
    disc_lr::Float64 = 0.0002
    gen_lr::Float64 = 0.0002
    device::Function = gpu
end

function load_MNIST_images(hparams)
    images = MNIST.traintensor(Float32)

    # Нормализуем изображения к (-1, 1)
    normalized_images = @. 2f0 * images - 1f0
    image_tensor = reshape(normalized_images, 28, 28, 1, :)

    # Создаем объект dataloader, который выполняет итерацию по мини-пакетам тензора изображений
    dataloader = DataLoader(image_tensor, batchsize=hparams.batch_size, shuffle=true)

    return dataloader
end

function create_output_image(gen, fixed_noise, hparams)
    fake_images = cpu(gen.(fixed_noise))
    image_array = reduce(vcat, reduce.(hcat, partition(fake_images, hparams.output_dim)))
    image_array = permutedims(dropdims(image_array; dims=(3, 4)), (2, 1))
    image_array = @. Gray(image_array + 1f0) / 2f0
    return image_array
end

function train_discriminator!(gen, disc, real_img, fake_img, opt, ps, hparams)

    disc_loss, grads = Flux.withgradient(ps) do
        discriminator_loss(disc(real_img), disc(fake_img))
    end

    # Обновляем параметры дискриминатора
    update!(opt, ps, grads)
    return disc_loss
end

function train_generator!(gen, disc, fake_img, opt, ps, hparams)

    gen_loss, grads = Flux.withgradient(ps) do
        generator_loss(disc(fake_img))
    end

    update!(opt, ps, grads)
    return gen_loss
end

function train(hparams)

    dev = hparams.device
    # Проверяем доступность CUDA
    if hparams.device == gpu
        if !CUDA.has_cuda()
        dev = cpu
        @warn "No gpu found, falling back to CPU"
        end
    end

    # Загружаем нормализованные изображения MNIST
    dataloader = load_MNIST_images(hparams)

    # Инициализируем модели и передаем их в соответствующее устройство
    disc = Discriminator() |> dev
    gen =  Generator(hparams.latent_dim) |> dev

    # Собираем параметры генератора и дискриминатора
    disc_ps = params(disc)
    gen_ps = params(gen)

    # Инициализируем оптимизаторы ADAM для обеих подмоделей
    # с соответствующими скоростями обучения
    disc_opt = ADAM(hparams.disc_lr)
    gen_opt = ADAM(hparams.gen_lr)

    # Создаем пакет фиксированного шума для визуализации обучения генератора с течением времени
    fixed_noise = [randn(Float32, hparams.latent_dim, 1) |> dev for _=1:hparams.output_dim^2]

    # Цикл обучения
    train_steps = 0
    for ep in 1:hparams.epochs
        @info "Epoch $ep"
        for real_img in dataloader

            # Передаем данные на GPU
            real_img = real_img |> dev

            # Создаем случайный шум
            noise = randn!(similar(real_img, (hparams.latent_dim, hparams.batch_size)))
            # Передаем шум в генератор для создания искусственного изображения
            fake_img = gen(noise)

            # Обновляем дискриминатор и генератор
            loss_disc = train_discriminator!(gen, disc, real_img, fake_img, disc_opt, disc_ps, hparams)
            loss_gen = train_generator!(gen, disc, fake_img, gen_opt, gen_ps, hparams)

            if train_steps % hparams.verbose_freq == 0
                @info("Train step $(train_steps), Discriminator loss = $(loss_disc), Generator loss = $(loss_gen)")
                # Сохраняем сгенерированное искусственное изображение
                output_image = create_output_image(gen, fixed_noise, hparams)
                save(@sprintf("output/dcgan_steps_%06d.png", train_steps), output_image)
            end
            train_steps += 1
        end
    end

    output_image = create_output_image(gen, fixed_noise, hparams)
    save(@sprintf("output/dcgan_steps_%06d.png", train_steps), output_image)

    return nothing
end

# Определяем гиперпараметры (здесь мы оставляем значения по умолчанию)
hparams = HyperParams()
train(hparams)

folder = "output"
# Получаем имена файлов изображений из папки
img_paths = readdir(folder, join=true)
# Загружаем все изображения как массив
images = load.(img_paths)
# Объединяем все изображения в массиве для создания матрицы изображений
gif_mat = cat(images..., dims=3)
save("./output.gif", gif_mat)