CODEWITHSUNDEEP
pythontensorflowpytorch
csobolew
csobolew

Reputation: 11

How can I use my existing TensorFlow weights in PyTorch?

I have a trained TensorFlow model that I am trying to convert to PyTorch. I recreated the neural network structure, and tried to manually copy the weights and biases for the layer from the h5 file, but I'm not getting the correct weights or prediction outputs when using the resultant PyTorch network.

I tried using only the first convolutional layer and copying the weights for that, but even still I am getting drastically different output from my TensorFlow code.

Here is an example of my model code for just the first convolutional layer in PyTorch:

class YourModel(nn.Module):
    def __init__(self):
        super(YourModel, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=6, kernel_size=7, stride=3)
        self.print = PrintLayer()

    def forward(self, x):
        x = self.conv1(x)
        x = nn.ReLU()(x)
        self.print(x)
        return x

and here is the TensorFlow model:

    def build_model(self):
        # Neural Net for Deep-Q learning Model
        model = Sequential()
        model.add(Conv2D(filters=6, kernel_size=(7, 7), strides=3, activation='relu', input_shape=(96, 96, self.frame_stack_num)))
        model.add(MaxPooling2D(pool_size=(2, 2)))
        model.add(Conv2D(filters=12, kernel_size=(4, 4), activation='relu'))
        model.add(MaxPooling2D(pool_size=(2, 2)))
        model.add(Flatten())
        model.add(Dense(216, activation='relu'))
        model.add(Dense(len(self.action_space), activation=None))
        model.compile(loss='mean_squared_error', optimizer=Adam(lr=self.learning_rate, epsilon=1e-7))
        return model

This is how I'm loading the weights into the PyTorch network:

tf_weights = h5py.File('trial_400.h5', 'r')
    set = tf_weights.keys()
    pytorch_model = YourModel()
    tmp = tf_weights['conv2d_3']
    pytorch_model.conv1.weight.data = torch.from_numpy(np.array(tf_weights['conv2d_3']["conv2d_3/kernel:0"]).transpose(3, 2, 1, 0)).float()
    pytorch_model.conv1.bias.data = torch.from_numpy(np.array(tf_weights['conv2d_3']["conv2d_3/bias:0"])).float()
    tf_weights.close()

Obviously there will be slight differences between the output numbers due to the randomization of the kernel, but I'm not getting remotely similar results. Any suggestions? Thanks.

Upvotes: 1

Views: 974

Answers (1)

Simon David
Simon David

Reputation: 1036

Luckily, there exist an established procedure to convert a TensorFlow model to PyTorch, which is why there is no need for re-creating the model and copy-pasting the weights.

PyTorch supports the Open Neural Network eXchange (ONNX) to represent machine-learning models. Therefore, you can export your trained TensorFlow model to ONNX using the tf2onnx converter, and subsequently load your ONNX model into PyTorch. Using these two steps should allow you to transform your TensorFlow model to PyTorch seamlessly.

Upvotes: 1

Related Questions