# Flux custom layers - incorrect solution from training

**URL:** https://discourse.julialang.org/t/flux-custom-layers-incorrect-solution-from-training/36993
**Category:** Machine Learning
**Created:** [April 3, 2020, 11:21pm UTC](https://discourse.julialang.org/t/flux-custom-layers-incorrect-solution-from-training/36993 "2020-04-03T23:21:20Z")
**Posts on this page:** 1
**Page:** 1

<div class="post-metadata">

### Author: ![lwhitefox](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/lwhitefox/32/4363_2.png) [@lwhitefox](https://discourse.julialang.org/u/lwhitefox)
#### Post date: [April 3, 2020, 11:21pm UTC](https://discourse.julialang.org/t/flux-custom-layers-incorrect-solution-from-training/36993/1 "2020-04-03T23:21:20Z")

</div>

My real goal is to include layers with matrices computed from known physics solutions; the training will then determine physical parameters appearing in the equations used to compute the matrix elements.

To get started, I created a toy problem with a 2x2 matrix. The equation generating matrix terms is:

```julia
function w(x::AbstractArray,y::Number)
  [3x[1]*y 2x[1]*x[1] ; -1.25y*y 2.5*x[1]*x[1]*y]
end

```

In the toy example, I only want to estimate the value of x, so it is set up as a 1 element vector, while y is left as a Number.

I compute a “true solution” (training data) using x=[1], y=2, and initialize a model with a value x\_guess=[1.2] as a single layer network, noting that the input to the network is a constant vector [1.;1.].

The model runs, the loss function works correctly, but the solution always converges to the same wrong answer, about 1.38.

The complete working example code is below, including some plotting lines that show that the error function has a minimum at x=1 as expected.

Does anyone have ideas why the network “training” does not get the right answer (i.e., spot the bug(s), sadly)? I’ve been staring at this but must be missing something; it should be able to get to zero error at x=1, but converges to the value x=1.38 with an error of 18.904071388502857.

```julia
using Flux

# trivial layer
struct mylayer{X <: AbstractArray, Y <: Number}
  x::X
  y::Y
end
# function that generates component of weight matrix
function w(x::AbstractArray,y::Number)
  [3x[1]*y 2x[1]*x[1] ; -1.25y*y 2.5*x[1]*x[1]*y]
end
function (a::mylayer)(vec::AbstractArray)
  x,y = a.x, a.y
  wmat = w(x,y)
  wmat*vec
end
# jazz it up with functor
Flux.@functor mylayer

# we want the true answer to be for x=y=1 - "input" is const [1.;1.]
y_true = w([1.0],2.0)*[1.0;1.0]

# now set up network configuration stuff
function loss(xx,yy)
  ŷ = m([1.0,1.0])
  Flux.mse(ŷ,yy)
end
opt = Descent(0.0005) # need to damp this or it blows up
# initialize model with [x]=1.2, y=2.0
x_guess = [1.02]
m = Chain(mylayer(x_guess,2.0))
# now train:
Flux.@epochs 15 Flux.train!(loss, params(m), zip(x_guess,y_true), opt, cb = () -> println("loss=$(loss([1.],y_true)) x=$(params(m))"))

# plot error function to show it is right, with minimum at [x]=[1.0]
using Plots
xvals = [x for x in 0:0.1:2.0]
err = [Flux.mse(w([x],2)*[1.;1.], y_true) for x in xvals]
Plots.plot(xvals,err)

```