Convert keras model to pytorch

Question

Is there an easy way to convert a model like this from keras to pytorch?

I have the code in keras as following:

from tensorflow.keras import Sequential
from tensorflow.keras.layers import Dense, Dropout
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.regularizers import l2

state_dim = 10
architecture = (256, 256)  # units per layer
learning_rate = 0.0001  # learning rate
l2_reg = 0.00000001  # L2 regularization
trainable = True
num_actions = 3

layers = []
n = len(architecture) # n = 2

for i, units in enumerate(architecture, 1):
    layers.append(Dense(units=units,
                        input_dim=state_dim if i == 1 else None,
                        activation='relu',
                        kernel_regularizer=l2(l2_reg),
                        name=f'Dense_{i}',
                        trainable=trainable))
    
layers.append(Dropout(.1))
layers.append(Dense(units=num_actions,
                    trainable=trainable,
                    name='Output'))

model = Sequential(layers)
model.compile(loss='mean_squared_error',
            optimizer=Adam(lr=learning_rate))

Which outputs as follow:

Model: "sequential_2"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
Dense_1 (Dense)              (None, 256)               2816      
_________________________________________________________________
Dense_2 (Dense)              (None, 256)               65792     
_________________________________________________________________
dropout_3 (Dropout)          (None, 256)               0         
_________________________________________________________________
Output (Dense)               (None, 3)                 771       
=================================================================
Total params: 69,379
Trainable params: 69,379
Non-trainable params: 0
_________________________________________________________________
None

I must admit, I'm a little out of my depth so any advice is appreciated. I'm trying to read through the pytorch docs and will update my question with a possible answer if I manage.

score 1 · Answer 1 · answered Jul 16 '21 at 19:05

Here is my best attempt:

state_dim = 10
architecture = (256, 256)  # units per layer
learning_rate = 0.0001  # learning rate
l2_reg = 0.00000001  # L2 regularization
trainable = True
num_actions = 3

import torch
from torch import nn

class CustomModel(nn.Module):
    def __init__(self):
        
        super().__init__()
        
        self.layers = nn.Sequential(
            nn.Linear(state_dim, architecture[0]),
            nn.ReLU(),
            nn.Linear(architecture[0], architecture[1]),
            nn.ReLU(),
            nn.Dropout(0.25),
            nn.Linear(architecture[1], num_actions),
        )

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

model = CustomModel()
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

It outputs a promising looking output:

CustomModel(
    (layers): Sequential(
        (0): Linear(in_features=10, out_features=256, bias=True)
        (1): ReLU()
        (2): Linear(in_features=256, out_features=256, bias=True)
        (3): ReLU()
        (4): Dropout(p=0.25, inplace=False)
        (5): Linear(in_features=256, out_features=3, bias=True)
    )
)

However a few items are still left unanswered:

are the activations in the right place?
how do we add a kernel_regularizer = l2(l2_reg) to the first two Linear/Dense layers?
and how do we make the layers trainable?

Any input appreciated.

Convert keras model to pytorch

1 Answers1