Initial commit

function_convertion_tests/test_bilinear_sampler.py

@@ -0,0 +1,44 @@
+import pickle
+import numpy as np
+import megengine as mge
+
+import torch
+import torch.nn.functional as F
+
+def bilinear_sampler(img, coords, mode='bilinear', mask=False):
+
+    """ Wrapper for grid_sample, uses pixel coordinates """
+    H, W = img.shape[-2:]
+    xgrid, ygrid = coords.split([1,1], dim=-1)
+    xgrid = 2*xgrid/(W-1) - 1
+    ygrid = 2*ygrid/(H-1) - 1
+
+    grid = torch.cat([xgrid, ygrid], dim=-1)
+    img = F.grid_sample(img, grid, align_corners=True)
+
+    if mask:
+        mask = (xgrid > -1) & (ygrid > -1) & (xgrid < 1) & (ygrid < 1)
+        return img, mask.float()
+
+    return img
+
+def test_bilinear_sampler():
+    # Getting back the megengine objects:
+    with open('test_data/bilinear_sampler_test.pickle', 'rb') as f:
+        right_feature_prev, coords, right_feature = pickle.load(f)
+
+    right_feature_prev = torch.tensor(right_feature_prev.numpy())
+    coords = torch.tensor(coords.numpy())
+    right_feature = right_feature.numpy()
+
+    # Test Pytorch
+    right_feature_pytorch = bilinear_sampler(right_feature_prev, coords).numpy()
+
+    error = np.mean(right_feature_pytorch-right_feature)
+    print(f"test_coords_grid - Avg. Error: {error},  \n \
+        Original shape: {coords.numpy().shape},\n  \
+        Obtained shape: {right_feature_pytorch.shape}, Expected shape: {right_feature.shape}")
+
+if __name__ == '__main__':
+
+    test_bilinear_sampler()

function_convertion_tests/test_coords_grid.py

@@ -0,0 +1,29 @@
+import pickle
+import numpy as np
+import megengine as mge
+
+import torch
+import torch.nn.functional as F
+
+def coords_grid(batch, ht, wd, device):
+    coords = torch.meshgrid(torch.arange(ht, device=device), torch.arange(wd, device=device), indexing='ij')
+    coords = torch.stack(coords[::-1], dim=0).float()
+    return coords[None].repeat(batch, 1, 1, 1)
+
+def test_coords_grid():
+    # Getting back the megengine objects:
+    with open('test_data/coords_grid_test.pickle', 'rb') as f:
+        batch, ht, wd, coords = pickle.load(f)
+
+    coords = coords.numpy()
+
+    # Test Pytorch
+    coords_pytorch = coords_grid(batch, ht, wd, 'cpu').numpy()
+
+    error = np.mean(coords_pytorch-coords)
+    print(f"test_coords_grid - Avg. Error: {error},  \n \
+        Obtained shape: {coords_pytorch.shape}, Expected shape: {coords.shape}")
+
+if __name__ == '__main__':
+
+    test_coords_grid()

function_convertion_tests/test_manual_pad.py

@@ -0,0 +1,51 @@
+import pickle
+import numpy as np
+import megengine as mge
+
+import torch
+import torch.nn.functional as F
+
+def manual_pad(x, pady, padx):
+
+    pad = (padx, padx, pady, pady)  
+    return F.pad(torch.tensor(x), pad, "replicate")
+
+
+def test_pad_1_1():
+    # Getting back the megengine objects:
+    with open('test_data/manual_pad_test1_1.pickle', 'rb') as f:
+        right_feature, pady, padx, right_pad = pickle.load(f)
+
+    right_feature = right_feature.numpy()
+    right_pad = right_pad.numpy()
+
+    # Test Pytorch
+    right_pad_pytorch = manual_pad(right_feature, pady, padx).numpy()
+
+    error = np.mean(right_pad_pytorch-right_pad)
+    print(f"test_pad_1_1 - Avg. Error: {error},  \n \
+        Orig. shape: {right_feature.shape}, \n \
+        Padded shape: {right_pad_pytorch.shape}, Expected shape: {right_pad.shape}")
+
+def test_pad_0_4():
+    # Getting back the megengine objects:
+    with open('test_data/manual_pad_test0_4.pickle', 'rb') as f:
+        right_feature, pady, padx, right_pad = pickle.load(f)
+
+    right_feature = right_feature.numpy()
+    right_pad = right_pad.numpy()
+
+    # Test Pytorch
+    right_pad_pytorch = manual_pad(right_feature, pady, padx).numpy()
+
+    error = np.mean(right_pad_pytorch-right_pad)
+    print(f"test_pad_0_4 - Avg. Error: {error},  \n \
+        Orig. shape: {right_feature.shape}, \n \
+        Padded shape: {right_pad_pytorch.shape}, Expected shape: {right_pad.shape}")
+
+
+if __name__ == '__main__':
+
+    test_pad_1_1()
+
+    test_pad_0_4()

function_convertion_tests/test_meshgrid.py

@@ -0,0 +1,30 @@
+import pickle
+import numpy as np
+import megengine as mge
+
+import torch
+import torch.nn.functional as F
+
+def test_meshgrid():
+    # Getting back the megengine objects:
+    with open('test_data/meshgrid_np_test.pkl', 'rb') as f:
+        rx, dilatex, ry, dilatey, x_grid, y_grid = pickle.load(f)
+
+    x_grid = x_grid.numpy()
+    y_grid = y_grid.numpy()
+
+    # Test Pytorch
+    x_grid_pytorch, y_grid_pytorch = torch.meshgrid(torch.arange(-rx, rx + 1, dilatex, device='cpu'), 
+                                     torch.arange(-ry, ry + 1, dilatey, device='cpu'), indexing='xy')
+
+
+    error_x = np.mean(x_grid_pytorch.numpy()-x_grid)
+    error_y = np.mean(y_grid_pytorch.numpy()-y_grid)
+    print(f"test_meshgrid (X) - Avg. Error: {error_x},  \n \
+        Obtained shape: {x_grid_pytorch.numpy().shape}, Expected shape: {x_grid.shape}")
+    print(f"test_meshgrid (Y) - Avg. Error: {error_y},  \n \
+        Obtained shape: {y_grid_pytorch.numpy().shape}, Expected shape: {y_grid.shape}")
+
+if __name__ == '__main__':
+
+    test_meshgrid()

function_convertion_tests/test_offset.py

@@ -0,0 +1,31 @@
+import pickle
+import numpy as np
+import megengine as mge
+
+import torch
+import torch.nn.functional as F
+
+def test_offset():
+    # Getting back the megengine objects:
+    with open('test_data/offset_test.pkl', 'rb') as f:
+        x_grid, y_grid, reshape_shape, transpose_order, expand_size, repeat_size, repeat_axis, offsets = pickle.load(f)
+
+    x_grid = torch.tensor(x_grid.numpy())
+    y_grid = torch.tensor(y_grid.numpy())
+    offsets_mge = offsets.numpy()
+    N = repeat_size
+
+    # Test Pytorch
+    offsets = torch.stack((x_grid, y_grid))
+    offsets = offsets.reshape(2, -1).permute(1, 0)
+    for d in sorted((0, 2, 3)):
+        offsets = offsets.unsqueeze(d)
+    offsets = offsets.repeat_interleave(N, dim=0)
+
+    error = np.mean(offsets.numpy()-offsets_mge)
+    print(f"test_offset - Avg. Error: {error},  \n \
+        Obtained shape: {offsets.numpy().shape}, Expected shape: {offsets_mge.shape}")
+
+if __name__ == '__main__':
+
+    test_offset()

function_convertion_tests/test_split.py

@@ -0,0 +1,47 @@
+import pickle
+import numpy as np
+import megengine as mge
+
+import torch
+import torch.nn.functional as F
+
+def test_split():
+    # Getting back the megengine objects:
+    with open('test_data/split_test.pkl', 'rb') as f:
+        left_feature, size, axis, lefts = pickle.load(f)
+
+    left_feature = torch.tensor(left_feature.numpy())
+
+    # Test Pytorch
+    lefts_pytorch = torch.split(left_feature, left_feature.shape[axis]//size, dim=axis)
+    
+    for i, (left_pytorch, left) in enumerate(zip(lefts_pytorch, lefts)):
+
+        error = np.mean(left_pytorch.numpy()-left.numpy())
+        print(f"test_split {i} - Avg. Error: {error},  \n \
+            Obtained shape: {left_pytorch.numpy().shape}, Expected shape: {left.numpy().shape}\n")
+
+def test_split_list():
+    # Getting back the megengine objects:
+    with open('test_data/split_test_list.pkl', 'rb') as f:
+        fmap1, size, axis, net, inp = pickle.load(f)
+
+    fmap1 = torch.tensor(fmap1.numpy())
+    net = net.numpy()
+    inp = inp.numpy()
+
+    # Test Pytorch
+    net_pytorch, inp_pytorch = torch.split(fmap1, [size[0],size[0]], dim=axis)
+
+    error_net = np.mean(net_pytorch.numpy()-net)
+    error_inp = np.mean(inp_pytorch.numpy()-inp)
+    print(f"test_split_list (net) - Avg. Error: {error_net},  \n \
+        Obtained shape: {net_pytorch.numpy().shape}, Expected shape: {net.shape}\n")
+    print(f"test_split_list (inp) - Avg. Error: {error_inp},  \n \
+        Obtained shape: {inp_pytorch.numpy().shape}, Expected shape: {inp.shape}\n")
+
+
+if __name__ == '__main__':
+
+    test_split()
+    test_split_list()