Merge pull request #22 from TwentyBN/3d_support

3d support

Author: gmalivenko

README.md

@@ -84,8 +84,10 @@ Layers:
 
 * Linear
 * Conv2d
+* Conv3d
 * ConvTranspose2d
 * MaxPool2d
+* MaxPool3d
 * AvgPool2d
 * Global average pooling (as special case of AdaptiveAvgPool2d)
 * Embedding

pytorch2keras/layers.py

@@ -28,12 +28,58 @@ def convert_conv(params, w_name, scope_name, inputs, layers, weights, short_name
         tf_name = 'C' + random_string(7)
     else:
         tf_name = w_name + str(random.random())
-    
+
     bias_name = '{0}.bias'.format(w_name)
     weights_name = '{0}.weight'.format(w_name)
     input_name = inputs[0]
 
-    if len(weights[weights_name].numpy().shape) == 4:
+    if len(weights[weights_name].numpy().shape) == 5: # 3D conv
+        W = weights[weights_name].numpy().transpose(2, 3, 4, 1, 0)
+        height, width, channels, n_layers, n_filters = W.shape
+        print(W.shape)
+
+        if bias_name in weights:
+            biases = weights[bias_name].numpy()
+            has_bias = True
+        else:
+            biases = None
+            has_bias = False
+
+        if params['pads'][0] > 0 or params['pads'][1] > 0:
+            padding_name = tf_name + '_pad'
+            padding_layer = keras.layers.ZeroPadding3D(
+                padding=(params['pads'][0],
+                         params['pads'][1],
+                         params['pads'][2]),
+                name=padding_name
+            )
+            layers[padding_name] = padding_layer(layers[input_name])
+            input_name = padding_name
+
+        weights = None
+        if has_bias:
+            weights = [W, biases]
+        else:
+            weights = [W]
+
+        print(len(weights), len(weights[0]), len(weights[0][0]),
+              len(weights[0][0][0]), len(weights[0][0][0][0]),
+              len(weights[0][0][0][0][0]))
+        conv = keras.layers.Conv3D(
+            filters=n_filters,
+            kernel_size=(channels, height, width),
+            strides=(params['strides'][0],
+                     params['strides'][1],
+                     params['strides'][2]),
+            padding='valid',
+            weights=weights,
+            use_bias=has_bias,
+            activation=None,
+            dilation_rate=params['dilations'][0],
+            name=tf_name
+        )
+        layers[scope_name] = conv(layers[input_name])
+    elif len(weights[weights_name].numpy().shape) == 4: # 2D conv
         W = weights[weights_name].numpy().transpose(2, 3, 1, 0)
         height, width, channels, n_filters = W.shape
 
@@ -71,7 +117,7 @@ def convert_conv(params, w_name, scope_name, inputs, layers, weights, short_name
             name=tf_name
         )
         layers[scope_name] = conv(layers[input_name])
-    else:
+    else: # 1D conv
         W = weights[weights_name].numpy().transpose(2, 1, 0)
         width, channels, n_filters = W.shape
 
@@ -333,6 +379,61 @@ def convert_maxpool(params, w_name, scope_name, inputs, layers, weights, short_n
     layers[scope_name] = pooling(layers[input_name])
 
 
+def convert_maxpool3(params, w_name, scope_name, inputs, layers, weights, short_names):
+    """
+    Convert 3d Max pooling.
+
+    Args:
+        params: dictionary with layer parameters
+        w_name: name prefix in state_dict
+        scope_name: pytorch scope name
+        inputs: pytorch node inputs
+        layers: dictionary with keras tensors
+        weights: pytorch state_dict
+        short_names: use short names for keras layers
+    """
+
+    print('Converting pooling ...')
+
+    if short_names:
+        tf_name = 'P' + random_string(7)
+    else:
+        tf_name = w_name + str(random.random())
+
+    if 'kernel_shape' in params:
+        height, width, depth = params['kernel_shape']
+    else:
+        height, width, depth = params['kernel_size']
+
+    if 'strides' in params:
+        stride_height, stride_width, stride_depth = params['strides']
+    else:
+        stride_height, stride_width, stride_depth = params['stride']
+    if 'pads' in params:
+        padding_h, padding_w, padding_d, _, _ = params['pads']
+    else:
+        padding_h, padding_w, padding_d = params['padding']
+    input_name = inputs[0]
+    if padding_h > 0 and padding_w > 0 and padding_d > 0:
+        padding_name = tf_name + '_pad'
+        padding_layer = keras.layers.ZeroPadding3D(
+            padding=(padding_h, padding_w, padding_d),
+            name=padding_name
+        )
+        layers[padding_name] = padding_layer(layers[inputs[0]])
+        input_name = padding_name
+
+    # Pooling type
+    pooling = keras.layers.MaxPooling3D(
+        pool_size=(height, width, depth),
+        strides=(stride_height, stride_width, stride_depth),
+        padding='valid',
+        name=tf_name
+    )
+
+    layers[scope_name] = pooling(layers[input_name])
+
+
 def convert_dropout(params, w_name, scope_name, inputs, layers, weights, short_names):
     """
     Convert dropout.
@@ -979,6 +1080,7 @@ def target_layer(x):
     'onnx::Gemm': convert_gemm,
     'onnx::MaxPool': convert_maxpool,
     'max_pool2d': convert_maxpool,
+    'aten::max_pool3d': convert_maxpool3,
     'onnx::AveragePool': convert_avgpool,
     'onnx::Dropout': convert_dropout,
     'onnx::BatchNormalization': convert_batchnorm,

tests/conv3d.py

@@ -0,0 +1,46 @@
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.autograd import Variable
+from pytorch2keras.converter import pytorch_to_keras
+
+
+class TestConv3d(nn.Module):
+    """Module for Conv2d conversion testing
+    """
+
+    def __init__(self, inp=10, out=16, kernel_size=3, bias=True):
+        super(TestConv3d, self).__init__()
+        self.conv3d = nn.Conv3d(inp, out, kernel_size=kernel_size, bias=bias)
+
+    def forward(self, x):
+        x = self.conv3d(x)
+        return x
+
+
+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, 30)
+        out = np.random.randint(1, 30)
+
+        model = TestConv3d(inp, out, kernel_size, inp % 2)
+
+        input_var = Variable(torch.randn(1, inp, inp, inp, inp))
+
+        output = model(input_var)
+
+        k_model = pytorch_to_keras(model,
+                                   input_var,
+                                   (inp, inp, inp, inp,),
+                                   verbose=True)
+
+        pytorch_output = output.data.numpy()
+        keras_output = k_model.predict(input_var.numpy())
+        error = np.max(pytorch_output - keras_output)
+        print("iteration: {}, error: {}".format(i, error))
+        if max_error < error:
+            max_error = error
+
+    print('Max error: {0}'.format(max_error))

tests/max_pool3d.py

@@ -0,0 +1,46 @@
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.autograd import Variable
+from pytorch2keras.converter import pytorch_to_keras
+
+
+class MaxPool(nn.Module):
+    """Module for MaxPool conversion testing
+    """
+
+    def __init__(self, inp=10, out=16, kernel_size=3, bias=True):
+        super(MaxPool, self).__init__()
+        self.conv3d = nn.Conv3d(inp, out, kernel_size=kernel_size, bias=bias)
+        self.pool3d = nn.MaxPool3d(kernel_size=3, padding=1)
+
+    def forward(self, x):
+        x = self.conv3d(x)
+        x = self.pool3d(x)
+        return x
+
+
+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, 30)
+        out = np.random.randint(1, 30)
+
+        model = MaxPool(inp, out, kernel_size, inp % 2)
+
+        input_np = np.random.uniform(0, 1, (1, inp, inp, inp, inp))
+        input_var = Variable(torch.FloatTensor(input_np))
+        output = model(input_var)
+
+        k_model = pytorch_to_keras(model, input_var, (inp, inp, inp, inp,), verbose=True)
+
+        pytorch_output = output.data.numpy()
+        keras_output = k_model.predict(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))