Updated converter in order to support the latest pytorch master.

Author: gmalivenko

pytorch2keras/converter.py

@@ -32,22 +32,23 @@ def set_training(model, mode):
             model.train(old_mode)
 
 
-def _optimize_trace(trace, aten):
+def _optimize_graph(graph, aten):
     # run dce first to eliminate dead parts of the graph that might have been
     # left behind by things like symbolic_override
-    torch._C._jit_pass_dce(trace)
-    torch._C._jit_pass_lint(trace)
-
-    torch._C._jit_pass_peephole(trace)
-    torch._C._jit_pass_lint(trace)
-    torch._C._jit_pass_onnx(trace, aten)
-    torch._C._jit_pass_lint(trace)
-    torch._C._jit_pass_onnx_peephole(trace)
-    torch._C._jit_pass_lint(trace)
-    torch._C._jit_pass_dce(trace)
-    torch._C._jit_pass_lint(trace)
-    torch._C._jit_pass_canonicalize(trace)
-    torch._C._jit_pass_lint(trace)
+    torch._C._jit_pass_dce(graph)
+    torch._C._jit_pass_lint(graph)
+
+    torch._C._jit_pass_peephole(graph)
+    torch._C._jit_pass_lint(graph)
+    graph = torch._C._jit_pass_onnx(graph, aten)
+    torch._C._jit_pass_lint(graph)
+    torch._C._jit_pass_onnx_peephole(graph)
+    torch._C._jit_pass_lint(graph)
+    torch._C._jit_pass_dce(graph)
+    torch._C._jit_pass_lint(graph)
+    graph = torch._C._jit_pass_canonicalize(graph)
+    torch._C._jit_pass_lint(graph)
+    return graph
 
 
 def get_node_id(node):
@@ -88,7 +89,8 @@ def pytorch_to_keras(
         raise RuntimeError("state_dict changed after running the tracer; "
                            "something weird is happening in your model!")
 
-    _optimize_trace(trace, False)
+    # _optimize_trace(trace, False)
+    trace.set_graph(_optimize_graph(trace.graph(), False))
 
     if verbose:
         print(trace.graph())

pytorch2keras/layers.py

@@ -738,29 +738,29 @@ def convert_upsample(params, w_name, scope_name, inputs, layers, weights):
 
 
 AVAILABLE_CONVERTERS = {
-    'Conv': convert_conv,
-    'ConvTranspose': convert_convtranspose,
-    'Flatten': convert_flatten,
-    'Gemm': convert_gemm,
-    'MaxPool': convert_maxpool,
+    'onnx::Conv': convert_conv,
+    'onnx::ConvTranspose': convert_convtranspose,
+    'onnx::Flatten': convert_flatten,
+    'onnx::Gemm': convert_gemm,
+    'onnx::MaxPool': convert_maxpool,
     'max_pool2d': convert_maxpool,
-    'AveragePool': convert_avgpool,
-    'Dropout': convert_dropout,
-    'BatchNormalization': convert_batchnorm,
-    'Add': convert_elementwise_add,
-    'Mul': convert_elementwise_mul,
-    'Sub': convert_elementwise_sub,
-    'Concat': convert_concat,
-    'Relu': convert_relu,
-    'LeakyRelu': convert_lrelu,
-    'Sigmoid': convert_sigmoid,
-    'Softmax': convert_softmax,
-    'Tanh': convert_tanh,
-    'Transpose': convert_transpose,
-    'Reshape': convert_reshape,
-    'MatMul': convert_matmul,
-    'Gather': convert_gather,
-    'ReduceSum': convert_reduce_sum,
-    'Constant': convert_constant,
-    'Upsample': convert_upsample,
+    'onnx::AveragePool': convert_avgpool,
+    'onnx::Dropout': convert_dropout,
+    'onnx::BatchNormalization': convert_batchnorm,
+    'onnx::Add': convert_elementwise_add,
+    'onnx::Mul': convert_elementwise_mul,
+    'onnx::Sub': convert_elementwise_sub,
+    'onnx::Concat': convert_concat,
+    'onnx::Relu': convert_relu,
+    'onnx::LeakyRelu': convert_lrelu,
+    'onnx::Sigmoid': convert_sigmoid,
+    'onnx::Softmax': convert_softmax,
+    'onnx::Tanh': convert_tanh,
+    'onnx::Transpose': convert_transpose,
+    'onnx::Reshape': convert_reshape,
+    'onnx::MatMul': convert_matmul,
+    'onnx::Gather': convert_gather,
+    'onnx::ReduceSum': convert_reduce_sum,
+    'onnx::Constant': convert_constant,
+    'onnx::Upsample': convert_upsample,
 }