Artificial Intelligence: Model migration from Keras to PyTorch (dense layer)

There are several alternatives to migrate keras.layers.Dense from Keras to Pytorch. In this particular case our keras.layers.Dense layer has a Softmax activation function.

Tensorflow/Keras

dense = keras.layers.Dense(2, activation='softmax')        

Pytorch

Pytorch doesn't include a default Dense class so you must import it from an external package or define it yourself.

Keras Dense has 2 steps: a linear function and activation function (can be Relu, Elu, Softmax)

# Get activation function or class
def get_activation(type):
  activation = None
  match type:
    case 'relu':
      activation = nn.ReLU()
    case 'elu':
      activation = nn.ELU()
    case 'softmax':
      activation = nn.Softmax(dim=1)
    case _: # Default model
      activation = f
  return activation

# A simple function to avoid forward conditions
# f(x) = x
# Can be replaced in forward by
# if self.activation: self.activation(x)
def f(x):
    return x

# CustomDense class
class CustomDense(nn.Module):
  def __init__(self, in, out, activation):
    super(CustomDense, self).__init__()
    self.linear = nn.Linear(in_features=in, out_features=out)
    self.activation = get_activation(activation)

  def forward(self, x):
    x = self.linear(x)
    x = self.activation(x)        

How to use it within your model.

self.dense = CustomDense(2, 2, 'softmax')