Документация Engee
Notebook

Генеративная сеть для синтезирования сигналов

В этом примере показано, как с помощью генеративной нейросети (GAN) можно синтезировать реалистичные радарные сигналы. Это может быть полезно для увеличения обучающей выборки, моделирования сценариев или генерации данных в условиях ограниченного доступа к реальным измерениям.

image_2.png

Генеративно-состязательная сеть — это две нейросети, в которой генератор берёт на вход шумовой сигнал и пытается из него сделать правдоподобный сигнал, а дискриминатор получает либо этот синтезированный сигнал, либо настоящий и учится отличать подделку от реальности, передавая генератору ошибку, чтобы тот со временем стал всё лучше маскировать свои фейки

Генератор получает на вход шум в виде одномерного сигнала из случайного распределения и превращает его через цепочку слоёв нейронной сети в синтезированный сигнал, который внешне должен стремиться быть похожим на настоящий сигнал

Дискриминатор получает на вход либо синтезированный сигнал от генератора либо реальный сигнал из тренировочного набора. После чего пропускает его через собственную сеть, чтобы принять решение, настоящее это или подделка, и выдаёт оценку вероятности

На основе этой оценки вычисляется ошибка для обеих частей системы и генератор учится улучшать свои синтезы, чтобы обманывать дискриминатор. Дискриминатор учится точнее отличать подделку от настоящего сигнала. По мере обучения обе ИНС улучшают свои способности и со временем синтезированный сигнал становится почти неотличим от реального

Импортируем необходимые пакеты для работы

In [ ]:
include("$(@__DIR__)/InstallPackages.jl")
In [ ]:
using Pkg
Pkg.instantiate()
using Glob
using DSP
using CUDA
using Flux 
using Statistics
using cuDNN
using BSON: @save, @load 
using Zygote

Поскольку используемые функции имеют достаточно большое количество строк кода, они были вынесены в отдельные файлы .jl, которые импортируются ниже Файл model.jl реализует структуру используемой модели. Файл dataset.jl используется для реализации логики модуля набора данных.

In [ ]:
Dirpath = "$(@__DIR__)"
include(joinpath(Dirpath, "model.jl"))
include(joinpath(Dirpath, "dataset.jl"))

Создание загрузчика данных

Далее мы создадим даталоадер, который будет батчами загружать данные в модель. В блоке ниже мы инициализуем датасет Логика датасета прописана в файле dataset.jl Датасет нарезает ЭКГ и радарный сигнал на окна с наложением. Нарезка позволяет снизить нагрузку на ИНС изза использования длинных последовательностей, а использование наложения позволяет учитывать временные зависимости

In [ ]:
function mycollate(batch)
    radar = hcat(first.(batch)...)           
    ecg   = hcat(last.(batch)...)          
    return radar, ecg
end
Out[0]:
mycollate (generic function with 1 method)
In [ ]:
ds = RadarECG("prepareRadarData");
N = length(ds)
train_idx = 1:N
X_train = [ ds[i][1] for i in train_idx ]  
Y_train = [ ds[i][2] for i in train_idx ]    
batch_size = 64
train_loader = Flux.DataLoader((X_train, Y_train);
    batchsize = batch_size,
    shuffle   = true,
    partial      = true,
    parallel = true,
    collate   = mycollate
);

После опредилим устройство, на котором будет обучаться модель. Если есть доступ к ГПУ, обучение будет ускорено на несколько порядков. Также перенесем загрузчик данных на используемый девайс

In [ ]:
device = CUDA.functional() ? gpu : cpu;
device == gpu ? train_loader = gpu.(train_loader) : nothing;

Инициализация моделей

Далее инициализируем ИНС - генератор и дискриминатор, а также определим испольщуемые оптимизаторы, в данном случае - Adam.

In [ ]:
nz = 100
net_G = Generator(nz)
net_D = Discriminator()
    
net_G = gpu(net_G)
net_D = gpu(net_D);

η_D, η_G = 5e-5, 2e-4
optD = Adam(η_D, (0.5, 0.999))            # без WeightDecay
optG = Adam(η_G, (0.5, 0.999));

Определим путь, куда будут сохраняться обученные модели

In [ ]:
isdir("output")||mkdir("output")
Dirpath = "$(@__DIR__)"
train_path = joinpath(Dirpath, "output");

Обучение модели

Определим функции, которые выполняют обучение GAN, а также функции потерь для каждой из подсетей - генератора и дискриминатора

In [ ]:
function discr_loss(real_output, fake_output)
    real_loss = Flux.logitbinarycrossentropy(real_output, 1f0)
    fake_loss = Flux.logitbinarycrossentropy(fake_output, 0f0)
    return (real_loss + fake_loss) / 2
end



generator_loss(fake_output) = Flux.logitbinarycrossentropy(fake_output, 1f0)

function train_discr(discr, original_data, fake_data, opt_discr)
    ps = Flux.params(discr)
    loss, back = Zygote.pullback(ps) do
                      discr_loss(discr(original_data), discr(fake_data))
    end
    grads = back(1f0)
    Flux.update!(opt_discr, ps, grads)
    return loss
end

Zygote.@nograd train_discr

function train_gan(gen, discr, original_data, opt_gen, opt_discr, nz, bsz)
    noise = randn(Float32, 2, nz, bsz) |> gpu
    loss = Dict()
    ps = Flux.params(gen)
    loss["gen"], back = Zygote.pullback(ps) do
                          fake_ = gen(noise)
                          loss["discr"] = train_discr(discr, original_data, fake_, opt_discr)
                          generator_loss(discr(fake_))
    end
    grads = back(1f0)
    Flux.update!(opt_gen, ps, grads)
    return loss
end
Out[0]:
train_gan (generic function with 1 method)

Запустим обучение

In [ ]:
train_steps = 0
epochs = 100
verbose_freq = 40
for ep in 1:epochs
    @info "Epoch $ep"
    for (i, (radar, _)) in enumerate(train_loader)
        radar = reshape(radar, size(radar,1), 1, size(radar,2)) |> gpu
        radar .+= 0.05f0 .* randn(Float32, size(radar)) |> gpu 
        loss = train_gan(net_G, net_D, radar, optG, optD, nz, batch_size)
        if train_steps % verbose_freq == 0
            noiseZ = randn(Float32,2,nz, batch_size) |> gpu
            P_real = mean(sigmoid.(net_D(radar)))
            P_fake = mean(sigmoid.(net_D(net_G(noiseZ))))
            @info("[$ep/$(epochs)] step $train_steps  " *
                  "P(real)=$(round(P_real,digits=3))  " *
                  "P(fake)=$(round(P_fake,digits=3))")
        end
        train_steps += 1
    end
end

Далее сохраним обученные модели

In [ ]:
# Переносим модели на CPU
train_dir = joinpath(Dirpath, "output")
net_D_cpu = cpu(net_D)
net_G_cpu = cpu(net_G)

# Пути для сохранения
output_dir_D = joinpath(train_dir, "modelsD_bson")
output_dir_G = joinpath(train_dir, "modelsG_bson")

# создаём папки, если их нет
isdir(output_dir_D) || mkdir(output_dir_D)
isdir(output_dir_G) || mkdir(output_dir_G)

# имена файлов для лучших моделей
best_path_D = joinpath(output_dir_D, "discriminator.bson")
best_path_G = joinpath(output_dir_G, "generator.bson")

# сохраняем лучший дискриминатор
@info "Сохраняем дискриминатор в $best_path_D"
@save best_path_D net_D_cpu

# сохраняем лучший генератор
@info "Сохраняем генератор в $best_path_G"
@save best_path_G net_G_cpu
[ Info: Сохраняем дискриминатор в /user/Диплом/GANECG/output/modelsD_bson/discriminator.bson
[ Info: Сохраняем генератор в /user/Диплом/GANECG/output/modelsG_bson/generator.bson

Тестирование модели

Сгенерируем из шума сигналы

In [ ]:
noise = randn(Float32, 1, 100, 16) |> gpu
fake  = net_G(noise)
fakec = cpu(fake);
ind = 2
plot(fakec[:, 1, ind]/6)
Out[0]:

Выводы

В данном демо примере была обучена нейросеть GAN для генерации синтетических данных на основе реальных