Autoencoder for telecommunication (Constellation shaping)

Specific Domains Machine Learning
1

Hi,

in telecommunication one can use a autoencoder-ish structure to learn a channel code/constellation shape. I’ve got a working Tensorflow implementation for constellation shaping, but cannot make it work in Julia. I think Flux somehow cannot handle the normalization I am trying to enforce. It just won’t train :frowning:

TF 1.14: https://github.com/Rassibassi/claude/blob/master/examples/tf_AutoEncoderForGeometricShapingAndAwgn.ipynb
Julia: https://github.com/Rassibassi/jaude/blob/master/ConstShapingAWGN.ipynb

or here:

using Flux
using Flux: @epochs, onehotbatch, throttle
using Statistics: mean, var, std
using Plots
pyplot()

add_dim(x::Array) = reshape(x, (1,size(x)...));

TR = Float32;
TC = ComplexF32;

M = 16;
constellation_dim = 2;
N = 32*M;
SNR = 20;
SNRlin = 10^(SNR/10) |> TR;

encoder = Chain(Dense(M, 32, Flux.relu), Dense(32, 32, Flux.relu), Dense(32, constellation_dim));
decoder = Chain(Dense(constellation_dim, 32, Flux.relu), Dense(32, 32, Flux.relu), Dense(32, M));

function model(X)
    X_seed = Flux.onehotbatch(1:M,1:M)
    s_seed = encoder(X_seed)
    s_seed = add_dim(s_seed[1,:] + 1im*s_seed[2,:])
    norm_factor = sqrt(mean(abs.(s_seed).^2))
    
    s = encoder(X)
    s = add_dim(s[1,:] + 1im*s[2,:]) / norm_factor
    𝜎 = sqrt(1/SNRlin) |> TR
    
    r = s + 𝜎 * randn(TC, 1, N)
    r = [real(r); imag(r)]
    Y = decoder(r)
    return Y
end    

loss(X) = Flux.logitcrossentropy(model(X), X);

opt = ADAM(0.001);
ps = params(encoder, decoder);

X = Flux.onehotbatch(rand(1:M, N), 1:M)
data = [[X]]

evalcb() = @show(loss(X));
@epochs 2000 Flux.train!(loss, ps, data, opt, cb = throttle(evalcb, 5));

X_seed = Flux.onehotbatch(1:M,1:M)
s_seed = encoder(X_seed)
s_seed_cpx = add_dim(s_seed[1,:] + 1im*s_seed[2,:])
norm_factor = sqrt(mean(abs.(s_seed_cpx).^2))

s = encoder(X)
s = add_dim(s[1,:] + 1im*s[2,:]) / norm_factor
𝜎 = sqrt(1/SNRlin) |> TR

r = s + 𝜎 * randn(TC, 1, N)
r = [real(r); imag(r)]

mean(abs.(s).^2)

scatter(Flux.Tracker.data(s_seed[1,:]),Flux.Tracker.data(s_seed[2,:]), markershape = :hexagon)

scatter(Flux.Tracker.data(r[1,:]),Flux.Tracker.data(r[2,:]), markershape = :hexagon)

Any ideas what I am missing?

2

In what way does it not train? Are gradients 0? Does it diverge?

3

I think Flux introduced pretty substantial changes internally not too long ago that made training a lot harder to access then it was maybe ~6 months ago. I could tell you how to make this work a few versions ago…

My assumption is the parameters aren’t updating at all (last time I touched Flux I had the same issue). But yea - more details would help a lot.

One thing to check is, the params struct, does it contain all the params used in your model?

4

Hej again,

sorry there was a mistake in my code! Somehow a cell in the notebook got lost. Now it actually runs, I edited my first post, but my initial problem is still there.

When you check out the notebook following the link after “TF 1.14”, at the bottom of the page, there is a plot of a constellation showing points in a rather symmetrical order. That’s how it should look like.

When I train the Julia version, it seems like it trains and converges,

loss(X) = 0.32920945f0 (tracked)
[ Info: Epoch 1999
loss(X) = 0.3057474f0 (tracked)
[ Info: Epoch 2000
loss(X) = 0.2962677f0 (tracked)

but somehow the plot of the constellation is not what I expect. The constellation does not look symmetrical at all. First plot in this notebook.

5

Hi, after some time I made it work! Here the code:

using Flux, Zygote
using Flux: @nograd, @epochs, onehotbatch, throttle, params
using Statistics: mean, var, std
using Plots
# pyplot()

add_dim(x::Array) = reshape(x, (1,size(x)...))

TR = Float32
TC = ComplexF32

M = 64
constellation_dim = 2
N = 4 * M
SNR = 20
SNRlin = 10^(SNR / 10) |> TR;

encoder = Chain(Dense(M, 32, Flux.relu), Dense(32, 32, Flux.relu), Dense(32, constellation_dim))
decoder = Chain(Dense(constellation_dim, 32, Flux.relu), Dense(32, 32, Flux.relu), Dense(32, M))

function model(X)
    X_seed = Flux.onehotbatch(1:M,1:M)
    s_seed = encoder(X_seed)
    s_seed = add_dim(complex.(s_seed[1,:], s_seed[2,:]))
    norm_factor = sqrt(mean(abs.(s_seed).^2))
    s_seed = s_seed / norm_factor

    s = encoder(X)
    s = add_dim(complex.(s[1,:], s[2,:])) / norm_factor
    𝜎 = sqrt(1/SNRlin) |> TR

    r = s + 𝜎 * randn(TC, 1, N)
    r = [real(r); imag(r)]
    Y = decoder(r)
    return Y , s_seed
end

function loss(x)
    Y, s_seed = model(x)
    return Flux.logitcrossentropy(Y, x)
end

@nograd onehotbatch
X = onehotbatch(rand(1:M, N), 1:M)
@show loss(X)

opt = ADAM(0.001)
ps = params(encoder, decoder)
data = [[X]]

evalcb() = @show(loss(X));
@epochs 4000 Flux.train!(loss, ps, data, opt, cb = throttle(evalcb, 5));

Y, s_seed = model(X)

@show mean(abs.(s_seed).^2)

scatter(real(s_seed)[1,:], imag(s_seed)[1,:], aspect_ratio = :equal, markershape = :hexagon)
ylims!((-2,2))
xlims!((-2,2))

const

1 Like