Turn warnings in prototype datasets tests into errors (#5540)

* fix PCAM prototype dataset

* update Zip and Tar archive loader datapipes

* only fail on warnings from the pytorch ecosystem

* Revert "only fail on warnings from the pytorch ecosystem"

This reverts commit 2bf3aa6f2d875a4055f7f3ed0b468316fc60d4f4.

test/builtin_dataset_mocks.py

@@ -10,6 +10,7 @@ import lzma
 import pathlib
 import pickle
 import random
+import warnings
 import xml.etree.ElementTree as ET
 from collections import defaultdict, Counter
 
@@ -470,6 +471,9 @@ def imagenet(info, root, config):
         ]
         num_children = 1
         synsets.extend((0, "", "", "", num_children, [], 0, 0) for _ in range(5))
+        with warnings.catch_warnings():
+            # The warning is not for savemat, but rather for some internals savemet is using
+            warnings.filterwarnings("ignore", category=np.VisibleDeprecationWarning)
             savemat(data_root / "meta.mat", dict(synsets=synsets))
 
         make_tar(root, devkit_root.with_suffix(".tar.gz").name, compression="gz")

test/test_prototype_builtin_datasets.py

@@ -35,6 +35,7 @@ def test_coverage():
         )
 
 
+@pytest.mark.filterwarnings("error")
 class TestCommon:
     @parametrize_dataset_mocks(DATASET_MOCKS)
     def test_smoke(self, test_home, dataset_mock, config):

torchvision/prototype/datasets/_builtin/imagenet.py

@@ -10,7 +10,7 @@ from torchdata.datapipes.iter import (
     Mapper,
     Filter,
     Demultiplexer,
-    TarArchiveReader,
+    TarArchiveLoader,
     Enumerator,
 )
 from torchvision.prototype.datasets.utils import (
@@ -158,7 +158,7 @@ class ImageNet(Dataset):
 
             # the train archive is a tar of tars
             if config.split == "train":
-                dp = TarArchiveReader(dp)
+                dp = TarArchiveLoader(dp)
 
             dp = hint_sharding(dp)
             dp = hint_shuffling(dp)

torchvision/prototype/datasets/_builtin/pcam.py

@@ -99,7 +99,7 @@ class PCAM(Dataset):
         image, target = data  # They're both numpy arrays at this point
 
         return {
-            "image": features.Image(image),
+            "image": features.Image(image.transpose(2, 0, 1)),
             "label": Label(target.item()),
         }
 

torchvision/prototype/datasets/utils/_resource.py

@@ -10,8 +10,8 @@ from torchdata.datapipes.iter import (
     FileLister,
     FileOpener,
     IterDataPipe,
-    ZipArchiveReader,
-    TarArchiveReader,
+    ZipArchiveLoader,
+    TarArchiveLoader,
     RarArchiveLoader,
 )
 from torchvision.datasets.utils import (
@@ -72,8 +72,8 @@ class OnlineResource(abc.ABC):
         return dp
 
     _ARCHIVE_LOADERS = {
-        ".tar": TarArchiveReader,
-        ".zip": ZipArchiveReader,
+        ".tar": TarArchiveLoader,
+        ".zip": ZipArchiveLoader,
         ".rar": RarArchiveLoader,
     }