Генерация сигналов с помощью GAN

using Flux
using Flux: @functor 

using Flux
using Flux: unsqueeze
using Flux: Conv, ConvTranspose, BatchNorm, MaxPool, leakyrelu, cat

function dcgan_init!(model::Chain)
    for layer in Flux.params(model)
        if isa(layer, Conv) || isa(layer, ConvTranspose)
            layer.weight .= 0.02f0 .* randn!(similar(layer.weight))



        elseif isa(layer, BatchNorm)

            layer.γ .= 1 .+ 0.02f0 .* randn!(similar(layer.γ))
            fill!(layer.β, 0f0)
        end
    end
    return model
end

function weights_init!(model)
    for layer in Flux.params(model)
        if isa(layer, Conv)
            glorot_uniform!(layer.weight, 0.0f0, 0.02f0)
        elseif isa(layer, ConvTranspose)
            glorot_uniform!(layer.weight, 0.0f0, 0.02f0)
        elseif isa(layer, BatchNorm)
            glorot_uniform!(layer.γ, 1.0f0, 0.02f0)
            fill!(layer.β, 0f0)
        end
    end
end

# ---------- 2) СЕТИ ----------------------------------------------------------
struct Generator
    model::Chain
end        
Flux.@layer Generator
(g::Generator)(x) = g.model(x)

function Generator(nz)
    Chain(
        ConvTranspose((314,), nz => 512; stride=1, pad=0, bias=false),
        BatchNorm(512), relu,

        ConvTranspose((4,), 512 => 256; stride=2, pad=1, bias=false),
        BatchNorm(256), relu,

        ConvTranspose((4,), 256 => 128; stride=2, pad=1, bias=false),
        BatchNorm(128), relu,

        ConvTranspose((4,), 128 => 64;  stride=2, pad=1, bias=false),
        BatchNorm(64),  relu,

        ConvTranspose((4,), 64  => 1;   stride=2, pad=1, bias=false)
    ) |> dcgan_init! |> Generator
end

struct Discriminator
     model::Chain
end     
Flux.@layer Discriminator
(d::Discriminator)(x) = d.model(x)

function Discriminator()
    Chain(
        Conv((4,), 1 => 64;  stride=2, pad=1, bias=false),
        x -> leakyrelu(x, 0.1f0),

        Conv((4,), 64 => 128; stride=2, pad=1, bias=false),
        BatchNorm(128), x -> leakyrelu(x, 0.1f0),

        Conv((4,), 128 => 256; stride=2, pad=1, bias=false),
        BatchNorm(256), x -> leakyrelu(x, 0.1f0),

        Conv((4,), 256 => 512; stride=2, pad=1, bias=false),
        BatchNorm(512), x -> leakyrelu(x, 0.1f0),

        # Dropout убираем — добавите позже, когда стабилизируется
        x -> mean(x, dims=1),      # global-avg-pool -> (1,512,B)
        Flux.flatten,                   # -> (512,B)
        Dense(512, 1; bias=false)  # логит, bias=False как в PyTorch
    ) |> dcgan_init! |> Discriminator
end