Added multiple inputs support.

Warning, this commit breaks all the tests and API!

Author: gmalivenko

pytorch2keras/converter.py

@@ -84,7 +84,7 @@ def pytorch_to_keras(
     orig_state_dict_keys = _unique_state_dict(model).keys()
 
     with set_training(model, training):
-        trace, torch_out = torch.jit.get_trace_graph(model, args)
+        trace, torch_out = torch.jit.get_trace_graph(model, tuple(args))
 
     if orig_state_dict_keys != _unique_state_dict(model).keys():
         raise RuntimeError("state_dict changed after running the tracer; "
@@ -117,12 +117,16 @@ def pytorch_to_keras(
     K.set_image_data_format('channels_first')
 
     layers = dict()
-    layers['input'] = keras.layers.InputLayer(
-        input_shape=input_shape, name='input'
-    ).output
+    keras_inputs = []
+    for i in range(len(args)):
+        layers['input{0}'.format(i)] = keras.layers.InputLayer(
+            input_shape=input_shape[i], name='input{0}'.format(i)
+        ).output
+        keras_inputs.append(layers['input{0}'.format(i)])
 
     outputs = []
 
+    input_index = 0
     for node in nodes:
         node_inputs = list(node.inputs())
         node_input_names = []
@@ -131,7 +135,8 @@ def pytorch_to_keras(
                 node_input_names.append(get_node_id(node_input.node()))
 
         if len(node_input_names) == 0:
-            node_input_names.append('input')
+            node_input_names.append('input{0}'.format(input_index))
+            input_index += 1
 
         node_type = node.kind()
         # print(dir(node))
@@ -168,7 +173,7 @@ def pytorch_to_keras(
         if node_id in graph_outputs:
             outputs.append(layers[node_id])
 
-    model = keras.models.Model(inputs=layers['input'], outputs=outputs)
+    model = keras.models.Model(inputs=keras_inputs, outputs=outputs)
 
     if change_ordering:
         import numpy as np

tests/multiple_inputs.py

@@ -0,0 +1,45 @@
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.autograd import Variable
+from pytorch2keras.converter import pytorch_to_keras
+
+
+class TestMultipleInputs(nn.Module):
+    """Module for Conv2d conversion testing
+    """
+
+    def __init__(self, inp=10, out=16, kernel_size=3, bias=True):
+        super(TestMultipleInputs, self).__init__()
+        self.conv2d = nn.Conv2d(inp, out, kernel_size=kernel_size, bias=bias)
+
+    def forward(self, x, y, z):
+        return self.conv2d(x) + self.conv2d(y) + self.conv2d(z)
+
+
+if __name__ == '__main__':
+    max_error = 0
+    for i in range(100):
+        kernel_size = np.random.randint(1, 7)
+        inp = np.random.randint(kernel_size + 1, 100)
+        out = np.random.randint(1, 100)
+
+        model = TestMultipleInputs(inp, out, kernel_size, inp % 2)
+
+        input_np = np.random.uniform(0, 1, (1, inp, inp, inp))
+        input_var = Variable(torch.FloatTensor(input_np))
+        input_var2 = Variable(torch.FloatTensor(input_np))
+        input_var3 = Variable(torch.FloatTensor(input_np))
+        output = model(input_var, input_var2, input_var3)
+
+        k_model = pytorch_to_keras(model, [input_var, input_var2, input_var3], [(inp, inp, inp,), (inp, inp, inp,), (inp, inp, inp,)], verbose=True)
+        k_model.summary()
+        pytorch_output = output.data.numpy()
+        keras_output = k_model.predict([input_np, input_np, input_np])
+
+        error = np.max(pytorch_output - keras_output)
+        print(error)
+        if max_error < error:
+            max_error = error
+
+    print('Max error: {0}'.format(max_error))