Use torch.testing.assert_close in datasets_video_utils.py (#3875)

Co-authored-by: Philip Meier <[email protected]>

Author: NicolasHug

test/test_datasets_video_utils.py

@@ -7,6 +7,7 @@
 from torchvision.datasets.video_utils import VideoClips, unfold
 
 from common_utils import get_tmp_dir
+from _assert_utils import assert_equal
 
 
 @contextlib.contextmanager
@@ -40,46 +41,46 @@ def test_unfold(self):
             [0, 1, 2],
             [3, 4, 5],
         ])
-        self.assertTrue(r.equal(expected))
+        assert_equal(r, expected, check_stride=False)
 
         r = unfold(a, 3, 2, 1)
         expected = torch.tensor([
             [0, 1, 2],
             [2, 3, 4],
             [4, 5, 6]
         ])
-        self.assertTrue(r.equal(expected))
+        assert_equal(r, expected, check_stride=False)
 
         r = unfold(a, 3, 2, 2)
         expected = torch.tensor([
             [0, 2, 4],
             [2, 4, 6],
         ])
-        self.assertTrue(r.equal(expected))
+        assert_equal(r, expected, check_stride=False)
 
     @unittest.skipIf(not io.video._av_available(), "this test requires av")
     def test_video_clips(self):
         with get_list_of_videos(num_videos=3) as video_list:
             video_clips = VideoClips(video_list, 5, 5, num_workers=2)
-            self.assertEqual(video_clips.num_clips(), 1 + 2 + 3)
+            assert video_clips.num_clips() == 1 + 2 + 3
             for i, (v_idx, c_idx) in enumerate([(0, 0), (1, 0), (1, 1), (2, 0), (2, 1), (2, 2)]):
                 video_idx, clip_idx = video_clips.get_clip_location(i)
-                self.assertEqual(video_idx, v_idx)
-                self.assertEqual(clip_idx, c_idx)
+                assert video_idx == v_idx
+                assert clip_idx == c_idx
 
             video_clips = VideoClips(video_list, 6, 6)
-            self.assertEqual(video_clips.num_clips(), 0 + 1 + 2)
+            assert video_clips.num_clips() == 0 + 1 + 2
             for i, (v_idx, c_idx) in enumerate([(1, 0), (2, 0), (2, 1)]):
                 video_idx, clip_idx = video_clips.get_clip_location(i)
-                self.assertEqual(video_idx, v_idx)
-                self.assertEqual(clip_idx, c_idx)
+                assert video_idx == v_idx
+                assert clip_idx == c_idx
 
             video_clips = VideoClips(video_list, 6, 1)
-            self.assertEqual(video_clips.num_clips(), 0 + (10 - 6 + 1) + (15 - 6 + 1))
+            assert video_clips.num_clips() == 0 + (10 - 6 + 1) + (15 - 6 + 1)
             for i, v_idx, c_idx in [(0, 1, 0), (4, 1, 4), (5, 2, 0), (6, 2, 1)]:
                 video_idx, clip_idx = video_clips.get_clip_location(i)
-                self.assertEqual(video_idx, v_idx)
-                self.assertEqual(clip_idx, c_idx)
+                assert video_idx == v_idx
+                assert clip_idx == c_idx
 
     @unittest.skipIf(not io.video._av_available(), "this test requires av")
     def test_video_clips_custom_fps(self):
@@ -89,8 +90,8 @@ def test_video_clips_custom_fps(self):
                 video_clips = VideoClips(video_list, num_frames, num_frames, fps, num_workers=2)
                 for i in range(video_clips.num_clips()):
                     video, audio, info, video_idx = video_clips.get_clip(i)
-                    self.assertEqual(video.shape[0], num_frames)
-                    self.assertEqual(info["video_fps"], fps)
+                    assert video.shape[0] == num_frames
+                    assert info["video_fps"] == fps
                     # TODO add tests checking that the content is right
 
     def test_compute_clips_for_video(self):
@@ -103,9 +104,9 @@ def test_compute_clips_for_video(self):
         clips, idxs = VideoClips.compute_clips_for_video(video_pts, num_frames, num_frames,
                                                          orig_fps, new_fps)
         resampled_idxs = VideoClips._resample_video_idx(int(duration * new_fps), orig_fps, new_fps)
-        self.assertEqual(len(clips), 1)
-        self.assertTrue(clips.equal(idxs))
-        self.assertTrue(idxs[0].equal(resampled_idxs))
+        assert len(clips) == 1
+        assert_equal(clips, idxs)
+        assert_equal(idxs[0], resampled_idxs)
 
         # case 2: all frames appear only once
         num_frames = 4
@@ -115,9 +116,9 @@ def test_compute_clips_for_video(self):
         clips, idxs = VideoClips.compute_clips_for_video(video_pts, num_frames, num_frames,
                                                          orig_fps, new_fps)
         resampled_idxs = VideoClips._resample_video_idx(int(duration * new_fps), orig_fps, new_fps)
-        self.assertEqual(len(clips), 3)
-        self.assertTrue(clips.equal(idxs))
-        self.assertTrue(idxs.flatten().equal(resampled_idxs))
+        assert len(clips) == 3
+        assert_equal(clips, idxs)
+        assert_equal(idxs.flatten(), resampled_idxs)
 
         # case 3: frames aren't enough for a clip
         num_frames = 32
@@ -126,8 +127,8 @@ def test_compute_clips_for_video(self):
         with self.assertWarns(UserWarning):
             clips, idxs = VideoClips.compute_clips_for_video(video_pts, num_frames, num_frames,
                                                              orig_fps, new_fps)
-        self.assertEqual(len(clips), 0)
-        self.assertEqual(len(idxs), 0)
+        assert len(clips) == 0
+        assert len(idxs) == 0
 
 
 if __name__ == '__main__':