| | |
| |
|
| | import math |
| | import unittest |
| | from itertools import permutations |
| |
|
| | import mlx.core as mx |
| | import mlx_tests |
| | import numpy as np |
| |
|
| | try: |
| | import torch |
| | import torch.nn.functional as F |
| |
|
| | has_torch = True |
| | except ImportError as e: |
| | has_torch = False |
| |
|
| |
|
| | class TestConvTranspose(mlx_tests.MLXTestCase): |
| | @unittest.skipIf(not has_torch, "requires Torch") |
| | def test_torch_conv_transpose_1D(self): |
| | def run_conv_transpose_1D( |
| | N, |
| | C, |
| | O, |
| | iH, |
| | kH, |
| | stride, |
| | padding, |
| | output_padding=0, |
| | dilation=1, |
| | groups=1, |
| | dtype="float32", |
| | atol=1e-5, |
| | ): |
| | with self.subTest( |
| | dtype=dtype, |
| | N=N, |
| | C=C, |
| | O=O, |
| | iH=iH, |
| | kH=kH, |
| | stride=stride, |
| | padding=padding, |
| | dilation=dilation, |
| | groups=groups, |
| | ): |
| | np_dtype = getattr(np, dtype) |
| | np.random.seed(0) |
| | in_np = np.random.normal(0, 1.0 / C, (N, iH, C)).astype(np_dtype) |
| | wt_np = np.random.normal(0, 1.0 / C, (O, kH, int(C / groups))).astype( |
| | np_dtype |
| | ) |
| |
|
| | in_mx, wt_mx = map(mx.array, (in_np, wt_np)) |
| | in_pt = torch.from_numpy(in_np.transpose(0, 2, 1)) |
| | wt_pt = torch.from_numpy(wt_np.transpose(2, 0, 1)) |
| |
|
| | out_mx = mx.conv_transpose1d( |
| | in_mx, |
| | wt_mx, |
| | stride=stride, |
| | padding=padding, |
| | dilation=dilation, |
| | groups=groups, |
| | ) |
| | out_pt = torch.conv_transpose1d( |
| | in_pt, |
| | wt_pt, |
| | stride=stride, |
| | padding=padding, |
| | dilation=dilation, |
| | groups=groups, |
| | ) |
| | out_pt = torch.transpose(out_pt, 2, 1) |
| |
|
| | self.assertEqual(out_pt.shape, out_mx.shape) |
| | self.assertTrue(np.allclose(out_pt.numpy(), out_mx, atol=atol)) |
| |
|
| | for dtype in ("float32",): |
| | for N, C, O in ( |
| | (1, 1, 1), |
| | (1, 6, 1), |
| | (1, 1, 6), |
| | (4, 32, 64), |
| | ): |
| | for iH, kH, stride, padding in ( |
| | (1, 1, 1, 0), |
| | (3, 3, 1, 0), |
| | (31, 5, 5, 2), |
| | ): |
| | run_conv_transpose_1D(N, C, O, iH, kH, stride, padding, dtype=dtype) |
| |
|
| | |
| | N, C, O = (4, 32, 64) |
| | for iH, kH, stride, padding in ( |
| | (1, 1, 1, 0), |
| | (3, 3, 1, 0), |
| | (31, 5, 5, 2), |
| | ): |
| | for group in (1,): |
| | run_conv_transpose_1D( |
| | N, C, O, iH, kH, stride, padding, groups=group, dtype=dtype |
| | ) |
| |
|
| | |
| | for tpose_in, tpose_wt in ( |
| | ((0, 2, 1), (0, 1, 2)), |
| | ((0, 2, 1), (0, 2, 1)), |
| | ): |
| | with self.subTest(name="strided", tpose_in=tpose_in, tpose_wt=tpose_wt): |
| | in_np = np.random.normal(0, 1.0 / 16, (16, 16, 16)).astype(np.float32) |
| | wt_np = np.random.normal(0, 1.0 / 16, (16, 16, 16)).astype(np.float32) |
| |
|
| | in_mx, wt_mx = map(mx.array, (in_np, wt_np)) |
| | in_mx_t = mx.transpose(in_mx, tpose_in) |
| | wt_mx_t = mx.transpose(wt_mx, tpose_wt) |
| | out_mx = mx.conv_transpose1d(in_mx_t, wt_mx_t) |
| |
|
| | in_pt = torch.from_numpy(in_np.transpose(tpose_in).transpose(0, 2, 1)) |
| | wt_pt = torch.from_numpy(wt_np.transpose(tpose_wt).transpose(2, 0, 1)) |
| |
|
| | out_pt = torch.conv_transpose1d(in_pt, wt_pt) |
| | out_pt = torch.transpose(out_pt, 2, 1) |
| |
|
| | self.assertEqual(out_pt.shape, out_mx.shape) |
| | self.assertTrue(np.allclose(out_pt.numpy(), out_mx, atol=1e-5)) |
| |
|
| | @unittest.skipIf(not has_torch, "requires Torch") |
| | def test_torch_conv_transpose_1D_grad(self): |
| | def run_conv_transpose1D_grad( |
| | N, |
| | C, |
| | O, |
| | iH, |
| | kH, |
| | stride, |
| | padding, |
| | dilation=1, |
| | groups=1, |
| | dtype="float32", |
| | atol=1e-5, |
| | ): |
| | with self.subTest( |
| | dtype=dtype, |
| | N=N, |
| | C=C, |
| | O=O, |
| | iH=iH, |
| | kH=kH, |
| | stride=stride, |
| | padding=padding, |
| | dilation=dilation, |
| | groups=groups, |
| | ): |
| | np_dtype = getattr(np, dtype) |
| | np.random.seed(0) |
| | |
| |
|
| | in_np = np.random.normal(0, 1.0 / C, (N, iH, C)).astype(np_dtype) |
| | wt_np = np.random.normal(0, 1.0 / C, (O, kH, C)).astype(np_dtype) |
| |
|
| | in_mx, wt_mx = map(mx.array, (in_np, wt_np)) |
| | in_pt = torch.from_numpy(in_np.transpose(0, 2, 1)).requires_grad_(True) |
| | wt_pt = torch.from_numpy(wt_np.transpose(2, 0, 1)).requires_grad_(True) |
| |
|
| | out_pt = F.conv_transpose1d( |
| | in_pt, wt_pt, stride=stride, padding=padding, dilation=dilation |
| | ) |
| |
|
| | |
| | out_pt.retain_grad() |
| | out_pt.sum().backward() |
| |
|
| | pt_grad_in = in_pt.grad.permute(0, 2, 1).numpy() |
| | pt_grad_wt = wt_pt.grad.permute(1, 2, 0).numpy() |
| |
|
| | ct_mx = mx.array(out_pt.grad.numpy().transpose(0, 2, 1)) |
| |
|
| | def f(a, b): |
| | return mx.conv_transpose1d( |
| | a, |
| | b, |
| | stride=stride, |
| | padding=padding, |
| | dilation=dilation, |
| | groups=groups, |
| | ) |
| |
|
| | _, outs_mx = mx.vjp( |
| | f, |
| | [ |
| | in_mx, |
| | wt_mx, |
| | ], |
| | [ |
| | ct_mx, |
| | ], |
| | ) |
| |
|
| | mx_grad_in, mx_grad_wt = outs_mx |
| |
|
| | self.assertEqual(pt_grad_in.shape, mx_grad_in.shape) |
| | self.assertEqual(in_mx.shape, mx_grad_in.shape) |
| | self.assertTrue(np.allclose(pt_grad_in, mx_grad_in, atol=atol)) |
| |
|
| | self.assertEqual(pt_grad_wt.shape, mx_grad_wt.shape) |
| | self.assertEqual(wt_mx.shape, mx_grad_wt.shape) |
| | self.assertTrue(np.allclose(pt_grad_wt, mx_grad_wt, atol=atol)) |
| |
|
| | for dtype in ("float32",): |
| | for N, C, O in ( |
| | (1, 1, 1), |
| | (1, 6, 1), |
| | (1, 1, 6), |
| | (4, 32, 64), |
| | ): |
| | for iH, kH, stride, padding in ( |
| | (1, 1, 1, 0), |
| | (3, 3, 1, 0), |
| | (31, 5, 5, 2), |
| | ): |
| | run_conv_transpose1D_grad( |
| | N, C, O, iH, kH, stride, padding, dtype=dtype |
| | ) |
| |
|
| | @unittest.skipIf(not has_torch, "requires Torch") |
| | def test_torch_conv_transpose_2D(self): |
| | def run_conv_transpose2D( |
| | N, |
| | C, |
| | O, |
| | idim, |
| | kdim, |
| | stride, |
| | padding, |
| | dilation=(1, 1), |
| | groups=1, |
| | dtype="float32", |
| | atol=1e-5, |
| | ): |
| | with self.subTest( |
| | dtype=dtype, |
| | N=N, |
| | C=C, |
| | O=O, |
| | idim=idim, |
| | kdim=kdim, |
| | stride=stride, |
| | padding=padding, |
| | dilation=dilation, |
| | groups=groups, |
| | ): |
| | np_dtype = getattr(np, dtype) |
| | np.random.seed(0) |
| | iH, iW = idim |
| | kH, kW = kdim |
| | scale = 1.0 / math.sqrt(kH * kW * C) |
| | in_np = np.random.normal(0.0, scale, (N, iH, iW, C)).astype(np_dtype) |
| | wt_np = np.random.normal(0.0, 1.0, (O, kH, kW, int(C / groups))).astype( |
| | np_dtype |
| | ) |
| |
|
| | in_mx, wt_mx = map(mx.array, (in_np, wt_np)) |
| | in_pt = torch.from_numpy(in_np.transpose(0, 3, 1, 2)).to("cpu") |
| | wt_pt = torch.from_numpy(wt_np.transpose(3, 0, 1, 2)).to("cpu") |
| |
|
| | out_mx = mx.conv_transpose2d( |
| | in_mx, |
| | wt_mx, |
| | stride=stride, |
| | padding=padding, |
| | dilation=dilation, |
| | groups=groups, |
| | ) |
| | out_pt = torch.conv_transpose2d( |
| | in_pt, |
| | wt_pt, |
| | stride=stride, |
| | padding=padding, |
| | dilation=dilation, |
| | groups=groups, |
| | ) |
| | out_pt = torch.permute(out_pt, (0, 2, 3, 1)).numpy(force=True) |
| |
|
| | self.assertEqual(out_pt.shape, out_mx.shape) |
| | self.assertTrue(np.allclose(out_pt, out_mx, atol=atol)) |
| |
|
| | for dtype in ("float32",): |
| | for N, C, O in ( |
| | (1, 1, 1), |
| | (1, 6, 1), |
| | (1, 1, 6), |
| | (4, 32, 64), |
| | ): |
| | for idim, kdim, stride, padding in ( |
| | ((1, 1), (1, 1), (1, 1), (0, 0)), |
| | ((3, 3), (3, 1), (1, 1), (0, 0)), |
| | ((31, 31), (5, 5), (5, 5), (2, 2)), |
| | ): |
| | run_conv_transpose2D( |
| | N, C, O, idim, kdim, stride, padding, dtype=dtype |
| | ) |
| |
|
| | |
| | N, C, O = (4, 32, 64) |
| | for idim, kdim, stride, padding in ( |
| | ((1, 1), (1, 1), (1, 1), (0, 0)), |
| | ((3, 3), (3, 1), (1, 1), (0, 0)), |
| | ((31, 31), (5, 5), (5, 5), (2, 2)), |
| | ): |
| | for group in (1,): |
| | run_conv_transpose2D( |
| | N, C, O, idim, kdim, stride, padding, groups=group, dtype=dtype |
| | ) |
| |
|
| | @unittest.skipIf(not has_torch, "requires Torch") |
| | def test_torch_conv_transpose_2D_grad(self): |
| | def run_conv_transpose2D_grad( |
| | N, |
| | C, |
| | O, |
| | idim, |
| | kdim, |
| | stride, |
| | padding, |
| | dilation=(1, 1), |
| | groups=1, |
| | dtype="float32", |
| | atol=1e-5, |
| | ): |
| | with self.subTest( |
| | dtype=dtype, |
| | N=N, |
| | C=C, |
| | O=O, |
| | idim=idim, |
| | kdim=kdim, |
| | stride=stride, |
| | padding=padding, |
| | dilation=dilation, |
| | groups=groups, |
| | ): |
| | np_dtype = getattr(np, dtype) |
| | np.random.seed(0) |
| | iH, iW = idim |
| | kH, kW = kdim |
| | scale = 1.0 / math.sqrt(kH * kW * C * O) |
| |
|
| | in_np = np.random.normal(0.0, scale, (N, iH, iW, C)).astype(np_dtype) |
| | wt_np = np.random.normal(0.0, scale, (O, kH, kW, C)).astype(np_dtype) |
| |
|
| | in_mx, wt_mx = map(mx.array, (in_np, wt_np)) |
| | in_pt = torch.from_numpy(in_np.transpose(0, 3, 1, 2)).requires_grad_( |
| | True |
| | ) |
| | wt_pt = torch.from_numpy(wt_np.transpose(3, 0, 1, 2)).requires_grad_( |
| | True |
| | ) |
| |
|
| | out_pt = F.conv_transpose2d( |
| | in_pt, wt_pt, stride=stride, padding=padding, dilation=dilation |
| | ) |
| |
|
| | |
| | out_pt.retain_grad() |
| | out_pt.sum().backward() |
| |
|
| | pt_grad_in = in_pt.grad.permute(0, 2, 3, 1).numpy() |
| | pt_grad_wt = wt_pt.grad.permute(1, 2, 3, 0).numpy() |
| |
|
| | ct_mx = mx.array(out_pt.grad.numpy().transpose(0, 2, 3, 1)) |
| |
|
| | def f(a, b): |
| | return mx.conv_transpose2d( |
| | a, |
| | b, |
| | stride=stride, |
| | padding=padding, |
| | dilation=dilation, |
| | groups=groups, |
| | ) |
| |
|
| | _, outs_mx = mx.vjp( |
| | f, |
| | [in_mx, wt_mx], |
| | [ct_mx], |
| | ) |
| |
|
| | mx_grad_in, mx_grad_wt = outs_mx |
| |
|
| | self.assertEqual(pt_grad_in.shape, mx_grad_in.shape) |
| | self.assertEqual(in_mx.shape, mx_grad_in.shape) |
| | self.assertTrue(np.allclose(pt_grad_in, mx_grad_in, atol=atol)) |
| |
|
| | self.assertEqual(pt_grad_wt.shape, mx_grad_wt.shape) |
| | self.assertEqual(wt_mx.shape, mx_grad_wt.shape) |
| | self.assertTrue(np.allclose(pt_grad_wt, mx_grad_wt, atol=atol)) |
| |
|
| | for dtype in ("float32",): |
| | for N, C, O in ((1, 1, 1), (1, 6, 1), (1, 1, 6), (4, 32, 64), (4, 16, 32)): |
| | for idim, kdim, stride, padding, dilation in ( |
| | ((1, 1), (1, 1), (1, 1), (0, 0), (1, 1)), |
| | ((3, 3), (3, 1), (1, 1), (0, 0), (1, 1)), |
| | ((31, 31), (5, 5), (5, 5), (2, 2), (1, 1)), |
| | ((32, 32), (3, 3), (2, 2), (1, 1), (1, 1)), |
| | ((31, 31), (5, 5), (5, 5), (2, 2), (3, 2)), |
| | ((32, 32), (3, 3), (2, 2), (1, 1), (3, 2)), |
| | ): |
| | run_conv_transpose2D_grad( |
| | N, C, O, idim, kdim, stride, padding, dilation, dtype=dtype |
| | ) |
| |
|
| | @unittest.skipIf(not has_torch, "requires Torch") |
| | def test_torch_conv_transpose_3D(self): |
| | def run_conv_transpose3D( |
| | N, |
| | C, |
| | O, |
| | idim, |
| | kdim, |
| | stride, |
| | padding, |
| | dilation=(1, 1, 1), |
| | groups=1, |
| | dtype="float32", |
| | atol=1e-5, |
| | ): |
| | with self.subTest( |
| | dtype=dtype, |
| | N=N, |
| | C=C, |
| | O=O, |
| | idim=idim, |
| | kdim=kdim, |
| | stride=stride, |
| | padding=padding, |
| | dilation=dilation, |
| | groups=groups, |
| | ): |
| | np_dtype = getattr(np, dtype) |
| | np.random.seed(0) |
| | iD, iH, iW = idim |
| | kD, kH, kW = kdim |
| | scale = 1.0 / math.sqrt(kD * kH * kW * C * O) |
| | in_np = np.random.normal(0.0, scale, (N, iD, iH, iW, C)).astype( |
| | np_dtype |
| | ) |
| | wt_np = np.random.normal(0.0, 1.0, (O, kD, kH, kW, C)).astype(np_dtype) |
| |
|
| | in_mx, wt_mx = map(mx.array, (in_np, wt_np)) |
| | in_pt = torch.from_numpy(in_np.transpose(0, 4, 1, 2, 3)) |
| | wt_pt = torch.from_numpy(wt_np.transpose(4, 0, 1, 2, 3)) |
| |
|
| | out_mx = mx.conv_transpose3d( |
| | in_mx, |
| | wt_mx, |
| | stride=stride, |
| | padding=padding, |
| | dilation=dilation, |
| | groups=groups, |
| | ) |
| | out_pt = torch.conv_transpose3d( |
| | in_pt, |
| | wt_pt, |
| | stride=stride, |
| | padding=padding, |
| | dilation=dilation, |
| | groups=groups, |
| | ) |
| | out_pt = torch.permute(out_pt, (0, 2, 3, 4, 1)).numpy(force=True) |
| |
|
| | self.assertEqual(out_pt.shape, out_mx.shape) |
| | self.assertTrue(np.allclose(out_pt, out_mx, atol=atol)) |
| |
|
| | for dtype in ("float32",): |
| | for N, C, O in ( |
| | (1, 1, 1), |
| | (1, 6, 1), |
| | (1, 1, 6), |
| | (2, 8, 16), |
| | ): |
| | for idim, kdim, stride, padding in ( |
| | ((1, 1, 1), (1, 1, 1), (1, 1, 1), (0, 0, 0)), |
| | ((3, 3, 3), (3, 1, 1), (1, 1, 1), (0, 0, 0)), |
| | ((15, 15, 15), (3, 3, 3), (3, 3, 3), (2, 2, 2)), |
| | ): |
| | run_conv_transpose3D( |
| | N, C, O, idim, kdim, stride, padding, dtype=dtype |
| | ) |
| |
|
| | @unittest.skipIf(not has_torch, "requires Torch") |
| | def test_torch_conv_transpose_3D_grad(self): |
| | def run_conv_transpose3D_grad( |
| | N, |
| | C, |
| | O, |
| | idim, |
| | kdim, |
| | stride, |
| | padding, |
| | dilation=(1, 1, 1), |
| | groups=1, |
| | dtype="float32", |
| | atol=1e-4, |
| | ): |
| | with self.subTest( |
| | dtype=dtype, |
| | N=N, |
| | C=C, |
| | O=O, |
| | idim=idim, |
| | kdim=kdim, |
| | stride=stride, |
| | padding=padding, |
| | dilation=dilation, |
| | groups=groups, |
| | ): |
| | np_dtype = getattr(np, dtype) |
| | np.random.seed(0) |
| | iD, iH, iW = idim |
| | kD, kH, kW = kdim |
| | scale = 1.0 / math.sqrt(kD * kH * kW * C * O) |
| |
|
| | in_np = np.random.normal(0.0, scale, (N, iD, iH, iW, C)).astype( |
| | np_dtype |
| | ) |
| | wt_np = np.random.normal(0.0, scale, (O, kD, kH, kW, C)).astype( |
| | np_dtype |
| | ) |
| |
|
| | in_mx, wt_mx = map(mx.array, (in_np, wt_np)) |
| | in_pt = torch.from_numpy(in_np.transpose(0, 4, 1, 2, 3)).requires_grad_( |
| | True |
| | ) |
| | wt_pt = torch.from_numpy(wt_np.transpose(4, 0, 1, 2, 3)).requires_grad_( |
| | True |
| | ) |
| |
|
| | out_pt = F.conv_transpose3d( |
| | in_pt, |
| | wt_pt, |
| | stride=stride, |
| | padding=padding, |
| | dilation=dilation, |
| | groups=groups, |
| | ) |
| |
|
| | |
| | out_pt.retain_grad() |
| | out_pt.sum().backward() |
| |
|
| | pt_grad_in = in_pt.grad.permute(0, 2, 3, 4, 1).numpy() |
| | pt_grad_wt = wt_pt.grad.permute(1, 2, 3, 4, 0).numpy() |
| |
|
| | ct_mx = mx.array(out_pt.grad.numpy().transpose(0, 2, 3, 4, 1)) |
| |
|
| | def f(a, b): |
| | return mx.conv_transpose3d( |
| | a, |
| | b, |
| | stride=stride, |
| | padding=padding, |
| | dilation=dilation, |
| | groups=groups, |
| | ) |
| |
|
| | _, outs_mx = mx.vjp( |
| | f, |
| | [in_mx, wt_mx], |
| | [ct_mx], |
| | ) |
| |
|
| | mx_grad_in, mx_grad_wt = outs_mx |
| |
|
| | self.assertEqual(pt_grad_in.shape, mx_grad_in.shape) |
| | self.assertEqual(in_mx.shape, mx_grad_in.shape) |
| | self.assertTrue(np.allclose(pt_grad_in, mx_grad_in, atol=atol)) |
| |
|
| | self.assertEqual(pt_grad_wt.shape, mx_grad_wt.shape) |
| | self.assertEqual(wt_mx.shape, mx_grad_wt.shape) |
| | self.assertTrue(np.allclose(pt_grad_wt, mx_grad_wt, atol=atol)) |
| |
|
| | for dtype in ("float32",): |
| | for N, C, O in ((1, 1, 1), (1, 6, 1), (1, 1, 6), (2, 4, 8), (2, 8, 16)): |
| | for idim, kdim, stride, padding, dilation in ( |
| | ((1, 1, 1), (1, 1, 1), (1, 1, 1), (0, 0, 0), (1, 1, 1)), |
| | ((3, 3, 3), (3, 1, 1), (1, 1, 1), (0, 0, 0), (1, 1, 1)), |
| | ((7, 7, 7), (5, 5, 5), (5, 5, 5), (2, 2, 2), (1, 1, 1)), |
| | ((8, 8, 8), (3, 3, 3), (2, 2, 2), (1, 1, 1), (1, 1, 1)), |
| | ((7, 7, 7), (5, 5, 5), (3, 3, 3), (2, 2, 2), (3, 2, 2)), |
| | ((8, 8, 8), (3, 3, 3), (2, 2, 2), (1, 1, 1), (3, 2, 2)), |
| | ): |
| | run_conv_transpose3D_grad( |
| | N, C, O, idim, kdim, stride, padding, dilation, dtype=dtype |
| | ) |
| |
|
| | @unittest.skipIf(not has_torch, "requires Torch") |
| | def test_torch_conv_tranpose_1d_output_padding(self): |
| | def run_conv_transpose_1d_output_padding( |
| | N, C, O, iH, kH, stride, padding, output_padding, dtype="float32", atol=1e-5 |
| | ): |
| | with self.subTest( |
| | dtype=dtype, |
| | N=N, |
| | C=C, |
| | O=O, |
| | iH=iH, |
| | kH=kH, |
| | stride=stride, |
| | padding=padding, |
| | output_padding=output_padding, |
| | ): |
| | np_dtype = getattr(np, dtype) |
| | np.random.seed(0) |
| | in_np = np.random.normal(0, 1.0 / C, (N, iH, C)).astype(np_dtype) |
| | wt_np = np.random.normal(0, 1.0 / C, (O, kH, C)).astype(np_dtype) |
| |
|
| | in_mx, wt_mx = map(mx.array, (in_np, wt_np)) |
| | in_pt = torch.from_numpy(in_np.transpose(0, 2, 1)) |
| | wt_pt = torch.from_numpy(wt_np.transpose(2, 0, 1)) |
| |
|
| | out_mx = mx.conv_transpose1d( |
| | in_mx, |
| | wt_mx, |
| | stride=stride, |
| | padding=padding, |
| | output_padding=output_padding, |
| | ) |
| |
|
| | out_pt = torch.conv_transpose1d( |
| | in_pt, |
| | wt_pt, |
| | stride=stride, |
| | padding=padding, |
| | output_padding=output_padding, |
| | ) |
| | out_pt = torch.transpose(out_pt, 2, 1) |
| |
|
| | self.assertEqual(out_pt.shape, out_mx.shape) |
| | self.assertTrue(np.allclose(out_pt.numpy(), out_mx, atol=atol)) |
| |
|
| | for dtype in ("float32",): |
| | for N, C, O in ((1, 1, 1), (1, 6, 1), (4, 32, 64)): |
| | for iH, kH, stride, padding, output_padding in ( |
| | (3, 2, 2, 0, 1), |
| | (5, 3, 2, 1, 0), |
| | (7, 4, 3, 1, 2), |
| | ): |
| | run_conv_transpose_1d_output_padding( |
| | N, C, O, iH, kH, stride, padding, output_padding, dtype=dtype |
| | ) |
| |
|
| | @unittest.skipIf(not has_torch, "requires Torch") |
| | def test_torch_conv_transpose_2d_output_padding(self): |
| | def run_conv_transpose_2d_output_padding( |
| | N, |
| | C, |
| | O, |
| | idim, |
| | kdim, |
| | stride, |
| | padding, |
| | output_padding, |
| | dtype="float32", |
| | atol=1e-5, |
| | ): |
| | with self.subTest( |
| | dtype=dtype, |
| | N=N, |
| | C=C, |
| | O=O, |
| | idim=idim, |
| | kdim=kdim, |
| | stride=stride, |
| | padding=padding, |
| | output_padding=output_padding, |
| | ): |
| | np_dtype = getattr(np, dtype) |
| | np.random.seed(0) |
| | iH, iW = idim |
| | kH, kW = kdim |
| | in_np = np.random.normal(0, 1.0 / C, (N, iH, iW, C)).astype(np_dtype) |
| | wt_np = np.random.normal(0, 1.0 / C, (O, kH, kW, C)).astype(np_dtype) |
| |
|
| | in_mx, wt_mx = map(mx.array, (in_np, wt_np)) |
| | in_pt = torch.from_numpy(in_np.transpose(0, 3, 1, 2)) |
| | wt_pt = torch.from_numpy(wt_np.transpose(3, 0, 1, 2)) |
| |
|
| | out_mx = mx.conv_transpose2d( |
| | in_mx, |
| | wt_mx, |
| | stride=stride, |
| | padding=padding, |
| | output_padding=output_padding, |
| | ) |
| |
|
| | out_pt = torch.conv_transpose2d( |
| | in_pt, |
| | wt_pt, |
| | stride=stride, |
| | padding=padding, |
| | output_padding=output_padding, |
| | ) |
| | out_pt = torch.permute(out_pt, (0, 2, 3, 1)).numpy(force=True) |
| |
|
| | self.assertEqual(out_pt.shape, out_mx.shape) |
| | self.assertTrue(np.allclose(out_pt, out_mx, atol=atol)) |
| |
|
| | for dtype in ("float32",): |
| | for N, C, O in ((1, 1, 1), (1, 6, 1), (4, 32, 64)): |
| | for idim, kdim, stride, padding, output_padding in ( |
| | ((3, 3), (2, 2), (2, 2), (0, 0), (1, 1)), |
| | ((5, 5), (3, 3), (2, 2), (1, 1), (0, 0)), |
| | ((7, 7), (4, 4), (3, 3), (1, 1), (2, 2)), |
| | ): |
| | run_conv_transpose_2d_output_padding( |
| | N, |
| | C, |
| | O, |
| | idim, |
| | kdim, |
| | stride, |
| | padding, |
| | output_padding, |
| | dtype=dtype, |
| | ) |
| |
|
| | @unittest.skipIf(not has_torch, "requires Torch") |
| | def test_torch_conv_transpose_3d_output_padding(self): |
| | def run_conv_transpose_3d_output_padding( |
| | N, |
| | C, |
| | O, |
| | idim, |
| | kdim, |
| | stride, |
| | padding, |
| | output_padding, |
| | dtype="float32", |
| | atol=1e-5, |
| | ): |
| | with self.subTest( |
| | dtype=dtype, |
| | N=N, |
| | C=C, |
| | O=O, |
| | idim=idim, |
| | kdim=kdim, |
| | stride=stride, |
| | padding=padding, |
| | output_padding=output_padding, |
| | ): |
| | np_dtype = getattr(np, dtype) |
| | np.random.seed(0) |
| | iD, iH, iW = idim |
| | kD, kH, kW = kdim |
| | in_np = np.random.normal(0, 1.0 / C, (N, iD, iH, iW, C)).astype( |
| | np_dtype |
| | ) |
| | wt_np = np.random.normal(0, 1.0 / C, (O, kD, kH, kW, C)).astype( |
| | np_dtype |
| | ) |
| |
|
| | in_mx, wt_mx = map(mx.array, (in_np, wt_np)) |
| | in_pt = torch.from_numpy(in_np.transpose(0, 4, 1, 2, 3)) |
| | wt_pt = torch.from_numpy(wt_np.transpose(4, 0, 1, 2, 3)) |
| |
|
| | out_mx = mx.conv_transpose3d( |
| | in_mx, |
| | wt_mx, |
| | stride=stride, |
| | padding=padding, |
| | output_padding=output_padding, |
| | ) |
| | out_pt = torch.conv_transpose3d( |
| | in_pt, |
| | wt_pt, |
| | stride=stride, |
| | padding=padding, |
| | output_padding=output_padding, |
| | ) |
| | out_pt = torch.permute(out_pt, (0, 2, 3, 4, 1)).numpy(force=True) |
| |
|
| | self.assertEqual(out_pt.shape, out_mx.shape) |
| | self.assertTrue(np.allclose(out_pt, out_mx, atol=atol)) |
| |
|
| | for dtype in ("float32",): |
| | for N, C, O in ((1, 1, 1), (1, 6, 1), (4, 32, 64)): |
| | for idim, kdim, stride, padding, output_padding in ( |
| | ((3, 3, 3), (2, 2, 2), (2, 2, 2), (0, 0, 0), (1, 1, 1)), |
| | ((5, 5, 5), (3, 3, 3), (2, 2, 2), (1, 1, 1), (0, 0, 0)), |
| | ((7, 7, 7), (4, 4, 4), (3, 3, 3), (1, 1, 1), (2, 2, 2)), |
| | ): |
| | run_conv_transpose_3d_output_padding( |
| | N, |
| | C, |
| | O, |
| | idim, |
| | kdim, |
| | stride, |
| | padding, |
| | output_padding, |
| | dtype=dtype, |
| | ) |
| |
|
| |
|
| | if __name__ == "__main__": |
| | mlx_tests.MLXTestRunner() |
| |
|
