(N, H, W, C, kH, kW, O, strides, padding, groups, np_dtype)
| 55 | |
| 56 | |
| 57 | def bench_shape(N, H, W, C, kH, kW, O, strides, padding, groups, np_dtype): |
| 58 | scale = 1.0 / math.sqrt(kH * kH * C) |
| 59 | a_np = np.random.uniform(0, 0.5, (N, H, W, C)).astype(np_dtype) |
| 60 | b_np = np.random.uniform(-scale, scale, (O, kH, kW, int(C / groups))).astype( |
| 61 | np_dtype |
| 62 | ) |
| 63 | |
| 64 | a_mx = mx.array(a_np) |
| 65 | b_mx = mx.array(b_np) |
| 66 | |
| 67 | a_pt = torch.from_numpy(a_np.transpose((0, 3, 1, 2))).to("mps") |
| 68 | b_pt = torch.from_numpy(b_np.transpose((3, 0, 1, 2))).to("mps") |
| 69 | |
| 70 | torch.mps.synchronize() |
| 71 | |
| 72 | f_mx = make_mx_conv_transpose_2D(strides, padding, groups) |
| 73 | f_pt = make_pt_conv_transpose_2D(strides, padding, groups) |
| 74 | |
| 75 | time_torch = bench(f_pt, a_pt, b_pt) |
| 76 | time_mlx = bench(f_mx, a_mx, b_mx) |
| 77 | |
| 78 | out_mx = mx.conv_transpose2d( |
| 79 | a_mx, b_mx, stride=strides, padding=padding, groups=groups |
| 80 | ) |
| 81 | out_pt = torch.conv_transpose2d( |
| 82 | a_pt.to("cpu"), b_pt.to("cpu"), stride=strides, padding=padding, groups=groups |
| 83 | ) |
| 84 | out_pt = torch.permute(out_pt, (0, 2, 3, 1)) |
| 85 | out_pt = out_pt.numpy(force=True) |
| 86 | |
| 87 | atol = 2e-5 if np_dtype == np.float32 else 1e-4 |
| 88 | |
| 89 | if not np.allclose(out_pt, out_mx, atol=atol): |
| 90 | print( |
| 91 | f"Failed at {(N, H, W, C)}, {(O, kH, kW, C)} [strides = {strides}, padding = {padding}, groups = {groups}] with max(|a - b|) = {np.max(np.abs(out_pt - out_mx))}" |
| 92 | ) |
| 93 | |
| 94 | return time_mlx, time_torch |
| 95 | |
| 96 | |
| 97 | if __name__ == "__main__": |
no test coverage detected