convert_keras_to_pytorch.py

from tensorflow.keras import models
from tensorflow.keras import layers
from tensorflow.keras.models import load_model
import torch
import torch.nn as nn
import numpy as np
import sys, re

## read command line arguments and set in and out file names
if len(sys.argv) == 1:
    raise Exception("First argument (Keras model h5 file path)")
elif len(sys.argv) == 2:
    #     fin  = './tf_model/p20_model00_BEST.h5'
    fin = sys.argv[1]
    fout = fin + '.converted.pt'
    fout2 = fin + '.converted.traced.pt'
elif len(sys.argv) == 3:
    fin = sys.argv[1]
    fout = sys.argv[2]
    fout2 = fout + '.traced.pt'
elif len(sys.argv) == 4:
    fin = sys.argv[1]
    fout = sys.argv[2]
    fout2 = sys.argv[3]

print('\n<<< In/Out Model filenames >>>')
print('[in]  Keras   model         : %s'%fin)
print('[out] Pytorch model         : %s'%fout)
print('[out] Pytorch model (traced): %s'%fout2)

## Convert keras to pyt
# Load Keras model
model_keras = models.load_model(fin, compile=False)

# Read Keras layer by laer and convert each to PyTorch layer
print('\n<<< Keras model architecture >>>')
pt_stack = []
for k, klayer in enumerate(model_keras.layers):
    layer_name = re.sub(r'[^a-zA-Z]', '', klayer.name).lower()
    print('Layer-%03d: %s'%(k, layer_name))
    
    if layer_name in ['dense']:
        # create a linear layer
        _n_input = klayer.input_shape[1] # 0th dim is batch dim
        _n_output = klayer.output_shape[1]
        this_layer = nn.Linear(_n_input, _n_output)
        
        # load weights
        _weight = klayer.weights[0]
        _bias   = klayer.bias
        this_layer.weight.data = torch.from_numpy(_weight.numpy().transpose()) # weights
        this_layer.bias.data = torch.from_numpy(_bias.numpy())
        
        # add as a linear layer
        pt_stack.append(this_layer)
        
        # check activation function is embedded as a part of a dense layer
        act_type = re.sub(r'[^a-zA-Z]', '', klayer.activation.__name__).lower()
        if not act_type == 'linear':
            if act_type == 'relu':        pt_stack.append(nn.ReLU())
            elif act_type == 'sigmoid':   pt_stack.append(nn.Sigmoid())
            elif act_type == 'leakyrelu':
                _alpha = klayer.alpha.item()
                pt_stack.append(nn.LeakyReLU(negative_slope=_alpha))
            else:                         raise Exception("Add more activation types!")
        
    elif layer_name in ['relu']:
        pt_stack.append(nn.ReLU())
        
    elif layer_name in ['leakyrelu']:
        _alpha = klayer.alpha.item()
        pt_stack.append(nn.LeakyReLU(negative_slope=_alpha))
        
    elif layer_name in ['sigmoid']:
        pt_stack.append(nn.Sigmoid())    
        
    elif layer_name in ['batchnormalization']:
        # create a batchnorm layer
        _n_input     = klayer.input_shape[1]
        _epsilon     = klayer.epsilon
        _momentum    = klayer.momentum        
        this_layer = nn.BatchNorm1d(_n_input, eps = _epsilon, momentum = 1 - _momentum)
                     # note that Keras and Pyt use different convention for momentum
                     # this need to be updated for 2D or 3D (for convolution layers) 
        # load weights
        _gamma       = klayer.gamma
        _beta        = klayer.beta
        _moving_mean = klayer.moving_mean
        _moving_variance = klayer.moving_variance  
        this_layer.weight.data = torch.from_numpy(_gamma.numpy())
        this_layer.bias.data   = torch.from_numpy(_beta.numpy())
        this_layer.running_mean.data   = torch.from_numpy(_moving_mean.numpy())
        this_layer.running_var.data   = torch.from_numpy(_moving_variance.numpy())
        
        # add as a linear layer
        pt_stack.append(this_layer)
        
    elif layer_name in ['dropout']:
        _rate = klayer.rate
        pt_stack.append(nn.Dropout(_rate)) # this need to be updated for 2D or 3D (for convolution layers)
    
    else:
        raise Exception("Add more layer types")
        
# Construct PyTorch model
model_pt = nn.Sequential(*pt_stack)
model_pt.eval() # This sets a pytorch model to a evaluation (inference) mode by instructing
                # 1) to zero drop out; 2) not updating moving mean and variance; etc.

## Save
# save pytorch model
torch.save(model_pt, fout)
# save traced model
rand_input = torch.randn(2, model_keras.layers[0].input_shape[1])
model_pt_traced = torch.jit.trace(model_pt, rand_input)
model_pt_traced.save(fout2)

## Sanity check
rand_input = np.random.randn(2, model_keras.layers[0].input_shape[1])
predicted_keras = model_keras.predict(rand_input)
predicted_pt = model_pt(torch.Tensor(rand_input))
predicted_pt2 = model_pt_traced(torch.Tensor(rand_input))

print('\n<<< Sanity check: inference from a random input vector >>>')
print('Input (generated by np.random.rand):')
print(rand_input)
print('Output (Keras):')
print(predicted_keras[0])
print('Output (PyTorch):')
print(predicted_pt.detach().numpy()[0])
print('Output (PyTorch, traced):')
print(predicted_pt2.detach().numpy()[0])