# Updating only a subset of parameters in a neural network using Flux

**URL:** <https://discourse.julialang.org/t/updating-only-a-subset-of-parameters-in-a-neural-network-using-flux/72312>\
**Category:** Machine Learning\
**Created:** [November 30, 2021, 4:38pm UTC](https://discourse.julialang.org/t/updating-only-a-subset-of-parameters-in-a-neural-network-using-flux/72312 "2021-11-30T16:38:01Z")\
**Posts on this page:** 3\
**Page:** 1

<div class="post-metadata">

**Author:** ![Handam](https://avatars.discourse-cdn.com/v4/letter/h/ccd318/32.png) [@Handam](https://discourse.julialang.org/u/Handam)\
**Post date:** [November 30, 2021, 4:38pm UTC](https://discourse.julialang.org/t/updating-only-a-subset-of-parameters-in-a-neural-network-using-flux/72312/1 "2021-11-30T16:38:01Z")

</div>

I am currently playing around with actor-critic methods. In my current problem, the actor has four outputs:

```julia
pol = Chain(Dense(6, 32, relu), Dense(32, 4))

```

However, depending on the state only a subset of the actions are legal. So for choosing an action I use:

```julia
x = get_featurevector(s, pos_a)
probs = softmax(pol(x)[pos_a])
action = sample(1:length(probs), Weights(probs))

```

where `s` is the current state and `pos_a` is a subset of `[1,2,3,4]` denoting the possible actions.  
Now, when I update the actor (`pol`), my idea would be to ignore the outputs and the corresponding weights & biases that the actor weren’t allowed to pick. I would like to do something like:

```julia
ps = params(pol[1], pol[2].weight[pos_a, :], pol[2].bias[pos_a, :])
gs = gradient(ps) do
	Flux.logitcrossentropy(pol(x), target)
end
update!(opt(lr), ps, gs)

```

Does anyone how I can make this code work properly?  
I hope it is clear what I want to achieve - I tried to keep the explanation as short as possible. If you need a more detailed explanation, let me know!

---

<div class="post-metadata">

**Author:** ![trahflow](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/trahflow/32/30585_2.png) [@trahflow](https://discourse.julialang.org/u/trahflow)\
**Post date:** [November 30, 2021, 5:59pm UTC](https://discourse.julialang.org/t/updating-only-a-subset-of-parameters-in-a-neural-network-using-flux/72312/2 "2021-11-30T17:59:20Z")

</div>

Just multiplying the output of the actor with a boolean mask is not an option?

```julia
pol(x) .* (1:4 .∈ Ref(pos_a))

```

---

<div class="post-metadata">

**Author:** ![Handam](https://avatars.discourse-cdn.com/v4/letter/h/ccd318/32.png) [@Handam](https://discourse.julialang.org/u/Handam)\
**Post date:** [December 1, 2021, 5:25pm UTC](https://discourse.julialang.org/t/updating-only-a-subset-of-parameters-in-a-neural-network-using-flux/72312/3 "2021-12-01T17:25:53Z")

</div>

Thank you once again trahflow!  
As I am using the logitcrossentropy, I used

```julia
pol(x) + mask

```

with

```julia
mask = zeros(4)
for i in 1:4
	if !(i in pos_a)
		mask[i] = -Inf
	end
end

```