How to put everything into a Sequential container

1

I have this code that works and does what I want it to do. However, I would also need the Sequential container to contain not only the linear layers, but also the one hot encoding. How could I achieve this? Thanks!

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

preprocessing_layers = {}

## Input
obs_space = torch.Tensor([0, 1,2,3])
obs_space = obs_space.to(torch.long)

## Sequential
layers = [torch.nn.Linear(obs_space.size(-1), 5)]
preprocessing_layers["direction"] = nn.Sequential(*layers)

## One hot
obs_space = F.one_hot(obs_space,4)
obs_space = obs_space.to(torch.float32)

## Forward
a = preprocessing_layers["direction"](obs_space)
2

You could use the one hot encoding as the forward of a module and then stick that into the sequential, so this might go in:

class OneHot(torch.nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.dim = dim
    def forward(self, x):
        return F.one_hot(x, self.dim)

Best regards

Thomas

1 Like