Neural Network Topology Optimization with Adam Optimizer

Description

This tutorial demonstrates neural network-parametrized topology optimization using Flux.jl’s Adam optimizer. Unlike the IPOPT-based approach, this method uses first-order gradient descent with adaptive learning rates. A continuation strategy progressively increases the SIMP penalty (1→3) and constraint aggregation weight (1→100) to converge to binary designs.

The network uses 6 hidden layers with sigmoid output for density parametrization.

Setup

using TopOpt, Zygote, Flux
using NNlib: leakyrelu
using WGLMakie
WGLMakie.activate!(; resize_to=:parent)
using Bonito
if haskey(ENV, "QUARTO_PROJECT_DIR")
    Bonito.Page(exportable=true, offline=true)
else
    Bonito.browser_display()
end
display_app(app) = display(app)

WGLMakie.activate!(; resize_to=:parent) selects the browser renderer and fills the Quarto output column. Bonito.Page(exportable=true, offline=true) embeds the assets needed by visualize(...; static=true) in the Quarto output, so the visualization does not require a running Julia process.

Problem Definition

E = 1.0 # Young's modulus
v = 0.3 # Poisson's ratio
f = 1.0 # downward force
els = (160, 40)

problem = PointLoadCantilever(els, (1.0, 1.0), E, v, f)

V = 0.5       # volume fraction
xmin = 1e-6   # minimum density
rmin = 3.0    # filter radius

# Continuation parameters: penalty ramps 1 → 3 over the epochs, and the
# constraint weight ramps from 1 to 100 so the volume constraint is
# progressively enforced.
p = 1.0           # initial SIMP penalty
delta_p = 0.1     # penalty increment per epoch
p_max = 3.0       # maximum penalty

alpha = 1.0       # initial constraint weight
delta_alpha = 10.0
alpha_max = 100   # maximum constraint weight

# Initial solver and compliance for C0 (normalization constant)
penalty = PowerPenaltyFun(p)
solver = FEASolver(DirectSolver, problem; xmin, penalty)
cheqfilter = DensityFilterFun(solver; rmin)
comp = ComplianceFun(solver)
volfrac = VolumeFun(solver)
C0 = comp(cheqfilter(PseudoDensities(fill(V, getncells(problem)))))

Neural Network

# 6-layer MLP with sigmoid output (scalar density in [0, 1])
m = 20
act = leakyrelu
nn = NeuralNetworkFun(
    Chain(
        Dense(2, m, act; init=Flux.glorot_normal),
        Dense(m, m, act; init=Flux.glorot_normal),
        Dense(m, m, act; init=Flux.glorot_normal),
        Dense(m, m, act; init=Flux.glorot_normal),
        Dense(m, m, act; init=Flux.glorot_normal),
        Dense(m, 1, sigmoid; init=Flux.glorot_normal),
    ),
    problem;
    scale=true,
)
w0 = nn.init_params

Optimizer Setup

# Adam with a moderate learning rate; gradient clipping avoids spikes.
alg = Flux.Optimise.Adam(0.01)
clip_alg = Flux.Optimise.ClipValue(1.0)
w = copy(w0)

# Heaviside projection with nonzero steepness pushes densities toward 0/1.
proj = HeavisideProjectionFun(2.0)

# Termination criteria
eps = Inf
eps_star = 0.05    # target intermediate-density fraction
maxiter = 50
constr_tol = 0.01

Density Functions

function todensities(w; filter=true)
    if filter
        PseudoDensities(proj.(cheqfilter(nn(NNParams(w))).x))
    else
        PseudoDensities(proj.(nn(NNParams(w)).x))
    end
end

Optimization Loop

epoch = 1
while true
    global epoch, p, alpha, eps, w
    epoch > maxiter && break
    eps < eps_star && break

    # Ramp penalty and constraint weight
    global p = min(p + delta_p, p_max)
    global alpha = min(alpha + delta_alpha, alpha_max)

    # Update penalty on the existing solver (mutates in-place, so the
    # compliance and volume functions pick up the new penalty automatically)
    setpenalty!(solver, p)

    obj = w -> comp(todensities(w; filter=true)) / C0
    constr = w -> volfrac(todensities(w; filter=false)) / V - 1
    combined_obj = w -> obj(w) + alpha * constr(w)^2

    # Adam step with gradient clipping
    Δ = Zygote.gradient(combined_obj, w)[1]
    Δ = Flux.Optimise.apply!(clip_alg, w, Δ)
    Flux.Optimise.update!(alg, w, Δ)

    # Check convergence
    x = nn(NNParams(w)).x
    global eps = sum(0.05 .< x .< 0.95) / length(x)
    violation = abs(constr(w))

    println("Epoch $epoch: p=$p, α=$alpha, obj=$(round(obj(w); digits=4)), constr=$(round(constr(w); digits=4)), eps=$(round(eps; digits=3))")
    global epoch += 1
end

┌ Warning: Layer with Float32 parameters got Float64 input.
│   The input will be converted, but any earlier layers may be very slow.
│   layer = Dense(2 => 20, leakyrelu)  # 60 parameters
│   summary(x) = "2-element Vector{Float64}"
└ @ Flux ~/.julia/packages/Flux/hrg9M/src/layers/stateless.jl:60
Epoch 1: p=1.1, α=11.0, obj=0.8015, constr=0.3164, eps=1.0
Epoch 2: p=1.2000000000000002, α=21.0, obj=0.936, constr=0.2164, eps=1.0
Epoch 3: p=1.3000000000000003, α=31.0, obj=1.1908, constr=0.0755, eps=1.0
Epoch 4: p=1.4000000000000004, α=41.0, obj=1.769, constr=-0.1159, eps=1.0
Epoch 5: p=1.5000000000000004, α=51.0, obj=1.8942, constr=-0.1126, eps=1.0
Epoch 6: p=1.6000000000000005, α=61.0, obj=1.7704, constr=-0.0416, eps=1.0
Epoch 7: p=1.7000000000000006, α=71.0, obj=1.6074, constr=0.0438, eps=1.0
Epoch 8: p=1.8000000000000007, α=81.0, obj=1.5587, constr=0.0907, eps=1.0
Epoch 9: p=1.9000000000000008, α=91.0, obj=1.6191, constr=0.0979, eps=1.0
Epoch 10: p=2.000000000000001, α=100.0, obj=1.772, constr=0.0777, eps=1.0
Epoch 11: p=2.100000000000001, α=100.0, obj=2.0238, constr=0.0382, eps=1.0
Epoch 12: p=2.200000000000001, α=100.0, obj=2.3606, constr=-0.0079, eps=1.0
Epoch 13: p=2.300000000000001, α=100.0, obj=2.6226, constr=-0.0306, eps=1.0
Epoch 14: p=2.4000000000000012, α=100.0, obj=2.7203, constr=-0.0239, eps=1.0
Epoch 15: p=2.5000000000000013, α=100.0, obj=2.6809, constr=0.0036, eps=1.0
Epoch 16: p=2.6000000000000014, α=100.0, obj=2.5909, constr=0.0388, eps=1.0
Epoch 17: p=2.7000000000000015, α=100.0, obj=2.5484, constr=0.0655, eps=1.0
Epoch 18: p=2.8000000000000016, α=100.0, obj=2.5955, constr=0.0762, eps=1.0
Epoch 19: p=2.9000000000000017, α=100.0, obj=2.7322, constr=0.0721, eps=1.0
Epoch 20: p=3.0, α=100.0, obj=2.9568, constr=0.0562, eps=1.0
Epoch 21: p=3.0, α=100.0, obj=3.0657, constr=0.0338, eps=1.0
Epoch 22: p=3.0, α=100.0, obj=3.1636, constr=0.0148, eps=1.0
Epoch 23: p=3.0, α=100.0, obj=3.1783, constr=0.0074, eps=1.0
Epoch 24: p=3.0, α=100.0, obj=3.087, constr=0.0134, eps=1.0
Epoch 25: p=3.0, α=100.0, obj=2.9341, constr=0.0285, eps=1.0
Epoch 26: p=3.0, α=100.0, obj=2.7773, constr=0.0458, eps=1.0
Epoch 27: p=3.0, α=100.0, obj=2.6743, constr=0.0568, eps=1.0
Epoch 28: p=3.0, α=100.0, obj=2.6459, constr=0.0575, eps=1.0
Epoch 29: p=3.0, α=100.0, obj=2.6925, constr=0.0483, eps=1.0
Epoch 30: p=3.0, α=100.0, obj=2.7931, constr=0.034, eps=1.0
Epoch 31: p=3.0, α=100.0, obj=2.8656, constr=0.0253, eps=1.0
Epoch 32: p=3.0, α=100.0, obj=2.8351, constr=0.0282, eps=1.0
Epoch 33: p=3.0, α=100.0, obj=2.7237, constr=0.04, eps=1.0
Epoch 34: p=3.0, α=100.0, obj=2.6183, constr=0.052, eps=1.0
Epoch 35: p=3.0, α=100.0, obj=2.5745, constr=0.0566, eps=1.0
Epoch 36: p=3.0, α=100.0, obj=2.6017, constr=0.0522, eps=1.0
Epoch 37: p=3.0, α=100.0, obj=2.6766, constr=0.0423, eps=1.0
Epoch 38: p=3.0, α=100.0, obj=2.7451, constr=0.0334, eps=1.0
Epoch 39: p=3.0, α=100.0, obj=2.7495, constr=0.0315, eps=1.0
Epoch 40: p=3.0, α=100.0, obj=2.6875, constr=0.0371, eps=1.0
Epoch 41: p=3.0, α=100.0, obj=2.6126, constr=0.0449, eps=1.0
Epoch 42: p=3.0, α=100.0, obj=2.5759, constr=0.0485, eps=1.0
Epoch 43: p=3.0, α=100.0, obj=2.59, constr=0.0454, eps=1.0
Epoch 44: p=3.0, α=100.0, obj=2.6361, constr=0.0378, eps=1.0
Epoch 45: p=3.0, α=100.0, obj=2.6742, constr=0.0313, eps=1.0
Epoch 46: p=3.0, α=100.0, obj=2.6595, constr=0.0314, eps=1.0
Epoch 47: p=3.0, α=100.0, obj=2.5973, constr=0.0378, eps=1.0
Epoch 48: p=3.0, α=100.0, obj=2.5392, constr=0.0439, eps=1.0
Epoch 49: p=3.0, α=100.0, obj=2.5331, constr=0.043, eps=1.0
Epoch 50: p=3.0, α=100.0, obj=2.5735, constr=0.0356, eps=1.0

Results

obj_final = w -> comp(todensities(w; filter=true)) / C0
constr_final = w -> volfrac(todensities(w; filter=false)) / V - 1
@show obj_final(w)
@show constr_final(w)
@show eps

┌ Warning: Layer with Float32 parameters got Float64 input.
│   The input will be converted, but any earlier layers may be very slow.
│   layer = Dense(2 => 20, leakyrelu)  # 60 parameters
│   summary(x) = "2-element Vector{Float64}"
└ @ Flux ~/.julia/packages/Flux/hrg9M/src/layers/stateless.jl:60
obj_final(w) = 2.5734588096471738
constr_final(w) = 0.03563831482902735
eps = 1.0

1.0

Visualization

topology = todensities(w; filter=true)
fig = visualize(problem; static=true, topology=topology)
display_app(fig)

┌ Warning: Layer with Float32 parameters got Float64 input.
│   The input will be converted, but any earlier layers may be very slow.
│   layer = Dense(2 => 20, leakyrelu)  # 60 parameters
│   summary(x) = "2-element Vector{Float64}"
└ @ Flux ~/.julia/packages/Flux/hrg9M/src/layers/stateless.jl:60

undeformed mesh

load arrows

support arrows

Figure 1: Adam-optimized neural network design (6-layer MLP)