# How to use Flux.train! to train custom layer?

**URL:** <https://discourse.julialang.org/t/how-to-use-flux-train-to-train-custom-layer/28293>\
**Category:** Machine Learning\
**Tags:** question\
**Created:** [September 2, 2019, 2:53pm UTC](https://discourse.julialang.org/t/how-to-use-flux-train-to-train-custom-layer/28293 "2019-09-02T14:53:23Z")\
**Posts on this page:** 3\
**Page:** 1

<div class="post-metadata">

**Author:** ![zhangqian](https://avatars.discourse-cdn.com/v4/letter/z/b4bc9f/32.png) [@zhangqian](https://discourse.julialang.org/u/zhangqian)\
**Post date:** [September 2, 2019, 2:53pm UTC](https://discourse.julialang.org/t/how-to-use-flux-train-to-train-custom-layer/28293/1 "2019-09-02T14:53:23Z")

</div>

Hi

I am trying to build a toy custom layer according to the instruction from `Building Layers` ([Basics · Flux](https://fluxml.ai/Flux.jl/stable/models/basics/#Building-Layers-1)) in [Basics · Flux](https://fluxml.ai/Flux.jl/stable/models/basics/), and train it with the `train!` method in Flux.jl. However, the loss does not decrease and the params do not seem to change. Here are my code and output.

```julia
function Linear(in, out)
	W = param(randn(out,in))
    x -> W*x
end

model=Linear(10,1)
loss(x, y) = Flux.mse(model(x), y)
opt = ADAM()
dataset = repeated((train_data, target),10)
evalcb = () -> @show(loss(train_data, target))
println(params(model))
Flux.train!(loss, params(model), dataset, opt, cb=evalcb)

```

The output is

```julia
Params([])
loss(train_data, target) = 39.910205426070895 (tracked)
loss(train_data, target) = 39.910205426070895 (tracked)
loss(train_data, target) = 39.910205426070895 (tracked)
loss(train_data, target) = 39.910205426070895 (tracked)
loss(train_data, target) = 39.910205426070895 (tracked)
loss(train_data, target) = 39.910205426070895 (tracked)
loss(train_data, target) = 39.910205426070895 (tracked)
loss(train_data, target) = 39.910205426070895 (tracked)
loss(train_data, target) = 39.910205426070895 (tracked)
loss(train_data, target) = 39.910205426070895 (tracked)

```

How can I use the `train!` method to train the model with custom layer? Thank you very much!

---

<div class="post-metadata">

**Author:** ![improbable22](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/improbable22/32/5464_2.png) [@improbable22](https://discourse.julialang.org/u/improbable22)\
**Post date:** [September 2, 2019, 4:38pm UTC](https://discourse.julialang.org/t/how-to-use-flux-train-to-train-custom-layer/28293/2 "2019-09-02T16:38:32Z")

</div>

This is the problem:

> [@zhangqian](#):
>
> Params()

Since `params(model)` doesn’t contain any parameters, none get updated.

You can access the paramter by `Linear(2,2).W`, just because it’s enclosed by the function. So I think this works:

```julia
Flux.train!(loss, Params([model.W]), dataset, opt, cb=evalcb)

```

But if you look at what Flux does, it usually makes a struct to hold the parameters, which is then made callable. And `@treelike` adds this struct to the list of things which `params` (and some other functions) understand.

---

<div class="post-metadata">

**Author:** ![zhangqian](https://avatars.discourse-cdn.com/v4/letter/z/b4bc9f/32.png) [@zhangqian](https://discourse.julialang.org/u/zhangqian)\
**Post date:** [September 2, 2019, 6:45pm UTC](https://discourse.julialang.org/t/how-to-use-flux-train-to-train-custom-layer/28293/3 "2019-09-02T18:45:13Z")

</div>

The

```julia
Flux.train!(loss, Params([model.W]), dataset, opt, cb=evalcb)

```

does not work for me. The `params[]` remains blank. It seems that we cannot access the attribute in a function. But maintaining a struct to hold parameters works. Now the `train!` will update the parameters and decrease the loss. Thank you very much!