How to make the equivalent of Tensorflow sequential model in PyTorch?

I have a simple shallow model in Tensorflow for multiclass classification.

model = tf.keras.models.Sequential([  

    tf.keras.layers.Dense(64, input_shape = (384,), activation = "relu"),
    tf.keras.layers.Dense(36, activation="softmax", use_bias = False)
                          ])    
        
model.summary()  

Model: "sequential_5"
_________________________________________________________________
 Layer (type)                Output Shape              Param #   
=================================================================
 dense_13 (Dense)            (None, 64)                24640     
                                                                 
 dense_14 (Dense)            (None, 36)                2304 

I want to replicate it in PyTorch where I am using softmax in the forward function.

# base model - it works but gives v low accuracy. need to add more layers.

class LR(torch.nn.Module):

    def __init__(self, n_features):
        super(LR, self).__init__()
        self.lr = torch.nn.Linear(n_features, 36)
        nn.ReLU()

        
    def forward(self, x):
        out = torch.softmax(self.lr(x),dim = 1,dtype=None)
        return out

# parameters
n_features = 384
n_classes = 36
optim = torch.optim.SGD(model.parameters(), lr=0.1)
criterion = torch.nn.CrossEntropyLoss()

accuracy_fn = Accuracy(task="multiclass",  num_classes=36)

I have tried modifying my base model like so and it gives me an error later on.

class LR(torch.nn.Module):

    def __init__(self, n_features):
        super(LR, self).__init__()
        self.lr = torch.nn.Linear(n_features, 64)        
        self.lr = nn.ReLU()
        self.lr = torch.nn.Linear(64, 36)
        
    def forward(self, x):
        out = torch.softmax(self.lr(x),dim = 1,dtype=None)
        return out

epochs = 15

def train(model, optim, criterion, x, y, epochs=epochs):
    for e in range(1, epochs + 1):
        optim.zero_grad()
        out = model(x)
        loss = criterion(out, y)
        loss.backward()
        optim.step()
        print(f"Loss at epoch {e}: {loss.data}")

    return model

model = LR(n_features)

model = train(model, optim, criterion, X_train, y_train)
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[175], line 18
     12  #       acc = accuracy_fn(out, y)  
     13  #       acc.backward()
     14  #       optim.step()
     15  #       print(f"Acc at epoch {e}: {acc.data}")
     16     return model
---> 18 model = train(model, optim, criterion, X_train, y_train)

Cell In[175], line 6, in train(model, optim, criterion, x, y, epochs)
      4 for e in range(1, epochs + 1):
      5     optim.zero_grad()
----> 6     out = model(x)
      7     loss = criterion(out, y)
      8     loss.backward()

File /usr/local/lib/python3.10/site-packages/torch/nn/modules/module.py:1501, in Module._call_impl(self, *args, **kwargs)
   1496 # If we don't have any hooks, we want to skip the rest of the logic in
   1497 # this function, and just call forward.
   1498 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1499         or _global_backward_pre_hooks or _global_backward_hooks
   1500         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1501     return forward_call(*args, **kwargs)
   1502 # Do not call functions when jit is used
   1503 full_backward_hooks, non_full_backward_hooks = [], []

Cell In[172], line 10, in LR.forward(self, x)
      9 def forward(self, x):
---> 10     out = torch.softmax(self.lr(x),dim = 1,dtype=None)
     11     return out

File /usr/local/lib/python3.10/site-packages/torch/nn/modules/module.py:1501, in Module._call_impl(self, *args, **kwargs)
   1496 # If we don't have any hooks, we want to skip the rest of the logic in
   1497 # this function, and just call forward.
   1498 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1499         or _global_backward_pre_hooks or _global_backward_hooks
   1500         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1501     return forward_call(*args, **kwargs)
   1502 # Do not call functions when jit is used
   1503 full_backward_hooks, non_full_backward_hooks = [], []

File /usr/local/lib/python3.10/site-packages/torch/nn/modules/linear.py:114, in Linear.forward(self, input)
    113 def forward(self, input: Tensor) -> Tensor:
--> 114     return F.linear(input, self.weight, self.bias)

RuntimeError: mat1 and mat2 shapes cannot be multiplied (155648x384 and 64x36)

my data shapes are:

X_train.shape
torch.Size([155648, 384]) 
y_train.shape 
torch.Size([155648]).

Can someone give me a hand with my class LR and train function?

In your Torch model, you have overwritten the self.lr attribute 2 times.

class LR(torch.nn.Module):
    def __init__(self, n_features):
        super(LR, self).__init__()
        self.lr = torch.nn.Linear(n_features, 64)        
        self.lr = nn.ReLU()                  # lr is over-written here
        self.lr = torch.nn.Linear(64, 36)    # and here

So you end up with a model that is just this:

class LR(torch.nn.Module):
    def __init__(self, n_features):
        super(LR, self).__init__()
        self.lr = torch.nn.Linear(64, 36)

If you want the equivalent of the Keras model, you can use:

from torch import nn

model = nn.Sequential(
    nn.Linear(384, 64),
    nn.ReLU()
    nn.Linear(64, 36, bias=False)
    nn.Softmax()
)

If you want to build a class for your model, you can use:

import torch.nn as nn

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.L1 = nn.Linear(384, 64),
        self.RLU = nn.ReLU()
        self.L2 = nn.Linear(64, 36, bias=False)

    def forward(self, x):
        x = self.L1(x)
        x = self.RLU(x)
        x = self.L2(x)
        return x

I left off the softmax because the CrossEntropyLoss applies a log-softmax on it's own. If you are making predictions, you can apply a softmax separately.

You could do (using the Sequential), the following:

class LR(torch.nn.Module):
    def __init__(self, n_features):
        super(LR, self).__init__()
        self.lr = torch.nn.Sequential(
            torch.nn.Linear(n_features, 64),        
            torch.nn.ReLU(),
            torch.nn.Linear(64, 36, bias=False),
            torch.nn.Softmax(dim=1),
        )

    def forward(self, x):
        return self.lr(x)

This could be used with, e.g.,

x = torch.ones(1, 1000)

lr = LR(x.shape[1])
lr(x)