# How to use Flux to write a Multi-Head Output Network

**URL:** <https://discourse.julialang.org/t/how-to-use-flux-to-write-a-multi-head-output-network/118605>\
**Category:** Machine Learning\
**Created:** [August 26, 2024, 2:13am UTC](https://discourse.julialang.org/t/how-to-use-flux-to-write-a-multi-head-output-network/118605 "2024-08-26T02:13:39Z")\
**Posts on this page:** 5\
**Page:** 1

<div class="post-metadata">

**Author:** ![ZC\_H](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/zc_h/32/207967_2.png) [@ZC\_H](https://discourse.julialang.org/u/ZC_H)\
**Post date:** [August 26, 2024, 2:13am UTC](https://discourse.julialang.org/t/how-to-use-flux-to-write-a-multi-head-output-network/118605/1 "2024-08-26T02:13:39Z")

</div>

I want to use julia.Flux to recurrent SAC algorithm. But I am facing a problem when I write the policy network modular because it is a multi-head output network, which receive a state matrix as input and output both expectation and variance. It can be easy to handle this in pytorch like this:

```julia
class PolicyNetwork(nn.Module):
    def __init__ (self, n_states, n_actions, action_bounds, n_hidden_filters=256):
        super(PolicyNetwork, self). __init__ ()
        self.n_states = n_states
        self.n_hidden_filters = n_hidden_filters
        self.n_actions = n_actions
        self.action_bounds = action_bounds

        self.hidden1 = nn.Linear(in_features=self.n_states, out_features=self.n_hidden_filters)
        init_weight(self.hidden1) # 初始化该神经层的参数
        self.hidden1.bias.data.zero_() # 将偏差置零
        self.hidden2 = nn.Linear(in_features=self.n_hidden_filters, out_features=self.n_hidden_filters)
        init_weight(self.hidden2)
        self.hidden2.bias.data.zero_()

        self.mu = nn.Linear(in_features=self.n_hidden_filters, out_features=self.n_actions)
        init_weight(self.mu, initializer="xavier uniform")
        self.mu.bias.data.zero_()

        self.log_std = nn.Linear(in_features=self.n_hidden_filters, out_features=self.n_actions)
        init_weight(self.log_std, initializer="xavier uniform")
        self.log_std.bias.data.zero_()

    def forward(self, states):
        x = F.relu(self.hidden1(states))
        x = F.relu(self.hidden2(x))

        mu = self.mu(x)
        log_std = self.log_std(x) # 给出方差
        return mu, log_std

```

but in julia, there is no class, so I have no idea how to deal with it. may anybody give me a code example using Flux? thanks very much!

---

<div class="post-metadata">

**Author:** ![ZC\_H](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/zc_h/32/207967_2.png) [@ZC\_H](https://discourse.julialang.org/u/ZC_H)\
**Post date:** [August 26, 2024, 2:22am UTC](https://discourse.julialang.org/t/how-to-use-flux-to-write-a-multi-head-output-network/118605/2 "2024-08-26T02:22:09Z")

</div>

I ask gpt about this, and it write a code like this for me:

```julia
# 创建策略函数神经网络
struct PolicyNetwork
    hidden1::Chain
    hidden2::Chain
    mu::Dense
    std::Dense
end

function PolicyNetwork(n_state::Int, n_action::Int, n_hidden::Int)
    hidden1 = Chain(
        Dense(n_state => n_hidden, init=Flux.glorot_uniform),
        BatchNorm(n_hidden)
        )
    hidden2 = Chain(
        Dense(n_hidden => n_hidden, init=Flux.glorot_uniform),
        BatchNorm(n_hidden)
    )
    mu = Dense(n_hidden => n_action, init=Flux.glorot_uniform)
    std = Dense(n_hidden => n_action, init=Flux.glorot_uniform)
    return PolicyNetwork(hidden1, hidden2, mu, std)
end

function (network::PolicyNetwork)(states)
    x = network.hidden1(states)
    x = network.hidden2(x)
    mu = network.mu(x)
    std = network.std(x)
    return mu, std
end

## 测试PolicyNetwork
network = PolicyNetwork(n_state, n_action, n_hidden)
x = [1;;]
mu, std = network(x)

```

it can run. Maybe this is a method. If you have a better method, I will be appretirate if you can tell me 😍

---

<div class="post-metadata">

**Author:** ![ZC\_H](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/zc_h/32/207967_2.png) [@ZC\_H](https://discourse.julialang.org/u/ZC_H)\
**Post date:** [August 26, 2024, 2:39am UTC](https://discourse.julialang.org/t/how-to-use-flux-to-write-a-multi-head-output-network/118605/3 "2024-08-26T02:39:05Z")

</div>

plus, it can be very easy to transfer this network to gpu, just code like this:

```julia
function PolicyNetwork(n_state::Int, n_action::Int, n_hidden::Int)
    hidden1 = Chain(
        Dense(n_state => n_hidden, init=Flux.glorot_uniform),
        BatchNorm(n_hidden)
        )|>gpu
    hidden2 = Chain(
        Dense(n_hidden => n_hidden, init=Flux.glorot_uniform),
        BatchNorm(n_hidden)
    )|>gpu
    mu = Dense(n_hidden => n_action, init=Flux.glorot_uniform)|>gpu
    std = Dense(n_hidden => n_action, init=Flux.glorot_uniform)|>gpu
    return PolicyNetwork(hidden1, hidden2, mu, std)
end

```

others all the same

---

<div class="post-metadata">

**Author:** ![mcabbott](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mcabbott/32/6603_2.png) [@mcabbott](https://discourse.julialang.org/u/mcabbott)\
**Post date:** [August 26, 2024, 3:14am UTC](https://discourse.julialang.org/t/how-to-use-flux-to-write-a-multi-head-output-network/118605/4 "2024-08-26T03:14:51Z")

</div>

What AI doesn’t seem to know is that you need to tell Flux to look for parameters inside, by making the layer with a macro. Otherwise it will not be able to train:

```julia
julia> Flux.setup(Adam(), PolicyNetwork(2, 3, 4))
┌ Warning: setup found no trainable parameters in this model
└ @ Optimisers ~/.julia/packages/Optimisers/yDIWk/src/interface.jl:32
()

julia> Flux.@layer PolicyNetwork # Defines methods for functor, show

julia> Flux.setup(Adam(), PolicyNetwork(2, 3, 4)) # now this sees parameters
(hidden1 = (layers = ((weight = Leaf(Adam(0.001, (0.9, 0.999), 1.0e-8), (Float32[0.0 0.0; 0.0 0.0;

```

After this, `PolicyNetwork(2, 3, 4) |> gpu` will move all the parameters, so your second definition should not be needed.

Note also that adding type parameters might be a good idea, for performance:

```julia
struct PolicyNetwork{A,B,C,D}
    hidden1::A
    hidden2::B
    mu::C
    std::D
end

```

---

<div class="post-metadata">

**Author:** ![ZC\_H](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/zc_h/32/207967_2.png) [@ZC\_H](https://discourse.julialang.org/u/ZC_H)\
**Post date:** [August 26, 2024, 3:30am UTC](https://discourse.julialang.org/t/how-to-use-flux-to-write-a-multi-head-output-network/118605/5 "2024-08-26T03:30:34Z")

</div>

Thank you very much, I was just doubting why it cannot be trained 😍