How can I separate the last layer of deep network in pytorch?

Question

My deep network is:

self.actor = nn.Sequential(
   nn.Linear(state_dim, 256),
   nn.Softplus(),
   nn.Linear(256, 256),
   nn.Softplus(),
   nn.Linear(256, action_dim),
   nn.Softplus())

Now, I would like the network to give two separate outputs, like this:

That's to say, only the last layer of the network is different and maybe the last layer has different activation functions. How should I change my code above?

Answer 1

The model you want to build is not sequential anymore, since there are two parallel branches at the end. You can keep the common trunk and separate with two additional separate layers. Something like:

class Model(nn.Module):
    def __init__(self):
        super.__init__()

        self.actor = nn.Sequential(
            nn.Linear(state_dim, 256),
            nn.Softplus(),
            nn.Linear(256, 256),
            nn.Softplus())
      
        self.outA = nn.Sequential(
            nn.Linear(256, action_dim),
            nn.Softplus())

        self.outB = nn.Sequential(
            nn.Linear(256, action_dim),
            nn.Softplus())
  
    def forward(self, x):
      features = self.actor(x)
      return self.outA(features), self.outB(features)