Search code examples
deep-learningpytorchreinforcement-learning

PyTorch - Neural Network - Output single scalar value

Let's say we have the following neural network in PyTorch

seq_model = nn.Sequential(
nn.Linear(1, 13),
nn.Tanh(),
nn.Linear(13, 1))

With the following input tensor

input = torch.tensor([1.0, 1.0, 5.0], dtype=torch.float32).unsqueeze(1)

I can run forward through the net and get

seq_model(input)

tensor([[-0.0165],
        [-0.0165],
        [-0.2289]], grad_fn=<TanhBackward0>)

Probably I also can get a single scalar value as an output, but I'm not sure how. Thank you. I'm trying to use such an network for reinforcment learning, and use it as an value function approximator for game board state evaluation.

Solution

  • The first dimension of input represents the number of observations in your minibatch (3), the second dimension represents instead the number of features (1).

    If you want to forward a single 3d input, the network must be modified (nn.Linear(1, 13) becomes nn.Linear(3, 13)), and you must remove unsqueeze(1) on input. Otherwise, you can merge the three outputs by using a loss to compute a single scalar from them.