| 20 | @unittest.skipIf(not has_torch, "requires Torch") |
| 21 | def test_torch_upsample(self): |
| 22 | def run_upsample( |
| 23 | N, |
| 24 | C, |
| 25 | idim, |
| 26 | scale_factor, |
| 27 | mode, |
| 28 | align_corner, |
| 29 | dtype="float32", |
| 30 | atol=1e-5, |
| 31 | ): |
| 32 | with self.subTest( |
| 33 | N=N, |
| 34 | C=C, |
| 35 | idim=idim, |
| 36 | scale_factor=scale_factor, |
| 37 | mode=mode, |
| 38 | align_corner=align_corner, |
| 39 | ): |
| 40 | np_dtype = getattr(np, dtype) |
| 41 | np.random.seed(0) |
| 42 | iH, iW = idim |
| 43 | in_np = np.random.normal(-1.0, 1.0, (N, iH, iW, C)).astype(np_dtype) |
| 44 | |
| 45 | in_mx = mx.array(in_np) |
| 46 | in_pt = torch.from_numpy(in_np.transpose(0, 3, 1, 2)).to("cpu") |
| 47 | |
| 48 | out_mx = nn.Upsample( |
| 49 | scale_factor=scale_factor, |
| 50 | mode=mode, |
| 51 | align_corners=align_corner, |
| 52 | )(in_mx) |
| 53 | mode_pt = { |
| 54 | "nearest": "nearest", |
| 55 | "linear": "bilinear", |
| 56 | "cubic": "bicubic", |
| 57 | }[mode] |
| 58 | out_pt = F.interpolate( |
| 59 | in_pt, |
| 60 | scale_factor=scale_factor, |
| 61 | mode=mode_pt, |
| 62 | align_corners=align_corner if mode != "nearest" else None, |
| 63 | ) |
| 64 | out_pt = torch.permute(out_pt, (0, 2, 3, 1)).numpy(force=True) |
| 65 | self.assertEqual(out_pt.shape, out_mx.shape) |
| 66 | self.assertTrue(np.allclose(out_pt, out_mx, atol=atol)) |
| 67 | |
| 68 | for dtype in ("float32",): |
| 69 | for N, C in ((1, 1), (2, 3)): |