How to initialize one hot encoding in a class in PyTorch

1

I’m trying to convert to PyTorch this TF code.

preprocessing_layers[scalar_name] = tf.keras.models.Sequential(
        [multigrid_networks.one_hot_layer(scalar_dim),
         tf.keras.layers.Dense(scalar_fc)])

This is how I would write it in PyTorch:

## Sequential
layers = [F.one_hot(scalar_dim), torch.nn.Linear(obs_space.size(-1), 5)]
preprocessing_layers[scalar_name] = nn.Sequential(*layers)

The problem is that if I do so, PyTorch is throwing this error at me:

TypeError: one_hot(): argument 'input' (position 1) must be Tensor, not int

My question is: how do I put hot encoding in a Sequential container in the init method of a class, without inserting the input (which is what I’ll do in the forward method)?

2

Try this:

scalar_dim = 3
num_classes = 10

F.one_hot(torch.LongTensor([scalar_dim]), num_classes=num_classes)

Output:
tensor([[0, 0, 0, 1, 0, 0, 0, 0, 0, 0]])
3

Thanks :slight_smile:

This is how I rewrote the code following your piece of advice:

import torch.nn.functional as F
import torch.nn as nn
import torch

preprocessing_layers = {}

## Input
obs_space = torch.LongTensor([0,1,2,3])

## Sequential
scalar_dim = 3
num_classes = 4

layers = [F.one_hot(torch.LongTensor([scalar_dim]), num_classes=num_classes), torch.nn.Linear(obs_space.size(-1), 5)]
preprocessing_layers["direction"] = nn.Sequential(*layers)

But now I get this error:

TypeError: torch.LongTensor is not a Module subclass
4

You’ll have to wrap it in a Module child class, like this:

class ToOneHot(nn.Module):
    def __init__(self, num_classes):
        super(ToOneHot, self).__init__()
    def forward(self, x):
        return F.one_hot(torch.LongTensor([x], num_classes=num_classes))

nn.Sequential(ToOneHot(num_classes), torch.nn.Linear(obs_space.size(-1), 5))

Output:
Sequential(
  (0): ToOneHot()
  (1): Linear(in_features=4, out_features=5, bias=True)
)
1 Like
5

Brilliant, thank you!