(N, H, W, C, kH, kW, O, strides, padding, groups, np_dtype)
| 48 | |
| 49 | |
| 50 | def bench_shape(N, H, W, C, kH, kW, O, strides, padding, groups, np_dtype): |
| 51 | scale = 1.0 / math.sqrt(kH * kH * C) |
| 52 | a_np = np.random.uniform(0, 0.5, (N, H, W, C)).astype(np_dtype) |
| 53 | b_np = np.random.uniform(-scale, scale, (O, kH, kW, int(C / groups))).astype( |
| 54 | np_dtype |
| 55 | ) |
| 56 | |
| 57 | a_mx = mx.array(a_np) |
| 58 | b_mx = mx.array(b_np) |
| 59 | |
| 60 | a_pt = torch.from_numpy(a_np.transpose((0, 3, 1, 2))).to("cpu") |
| 61 | b_pt = torch.from_numpy(b_np.transpose((0, 3, 1, 2))).to("cpu") |
| 62 | |
| 63 | f_mx = make_mx_conv_2D(strides, padding, groups) |
| 64 | f_pt = make_pt_conv_2D(strides, padding, groups) |
| 65 | |
| 66 | time_torch = bench(f_pt, a_pt, b_pt) |
| 67 | time_mlx = bench(f_mx, a_mx, b_mx) |
| 68 | |
| 69 | out_mx = mx.conv2d(a_mx, b_mx, stride=strides, padding=padding, groups=groups) |
| 70 | out_pt = torch.conv2d( |
| 71 | a_pt.to("cpu"), b_pt.to("cpu"), stride=strides, padding=padding, groups=groups |
| 72 | ) |
| 73 | out_pt = torch.permute(out_pt, (0, 2, 3, 1)) |
| 74 | out_pt = out_pt.numpy(force=True) |
| 75 | |
| 76 | atol = 2e-5 if np_dtype == np.float32 else 1e-4 |
| 77 | |
| 78 | if not np.allclose(out_pt, out_mx, atol=atol): |
| 79 | print( |
| 80 | 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))}" |
| 81 | ) |
| 82 | |
| 83 | return time_mlx, time_torch |
| 84 | |
| 85 | |
| 86 | if __name__ == "__main__": |
no test coverage detected