Merge branch 'main' into prototype/preprocessing_refs

Author: datumbox

.circleci/config.yml

@@ -277,7 +277,9 @@ jobs:
           command: pip install --user --progress-bar=off pytest pytest-mock
       - run:
           name: Run tests
-          command: pytest test/test_prototype_*.py
+          command: pytest --junitxml=test-results/junit.xml -v --durations 20 test/test_prototype_*.py
+      - store_test_results:
+          path: test-results
 
   binary_linux_wheel:
     <<: *binary_common

.circleci/config.yml.in

@@ -42,6 +42,10 @@ def train_one_epoch(model, criterion, optimizer, data_loader, device, epoch, arg
         else:
             loss = criterion(output, target)
             loss.backward()
+
+        if args.clip_grad_norm is not None:
+            nn.utils.clip_grad_norm_(utils.get_optimizer_params(optimizer), args.clip_grad_norm)
+
         optimizer.step()
 
         if model_ema and i % args.model_ema_steps == 0:
@@ -472,6 +476,7 @@ def get_args_parser(add_help=True):
     parser.add_argument(
         "--train-crop-size", default=224, type=int, help="the random crop size used for training (default: 224)"
     )
+    parser.add_argument("--clip-grad-norm", default=None, type=float, help="the maximum gradient norm (default None)")
 
     # Prototype models only
     parser.add_argument("--weights", default=None, type=str, help="the weights enum name to load")

references/classification/train.py

@@ -409,3 +409,11 @@ def reduce_across_processes(val):
     dist.barrier()
     dist.all_reduce(t)
     return t
+
+
+def get_optimizer_params(optimizer):
+    """Generator to iterate over all parameters in the optimizer param_groups."""
+
+    for group in optimizer.param_groups:
+        for p in group["params"]:
+            yield p

references/classification/utils.py

@@ -5,4 +5,5 @@
 from .imagenet import ImageNet
 from .mnist import MNIST, FashionMNIST, KMNIST, EMNIST, QMNIST
 from .sbd import SBD
+from .semeion import SEMEION
 from .voc import VOC

torchvision/prototype/datasets/_builtin/__init__.py

@@ -0,0 +1,69 @@
+import io
+from typing import Any, Callable, Dict, List, Optional, Tuple
+
+import torch
+from torchdata.datapipes.iter import (
+    IterDataPipe,
+    Mapper,
+    Shuffler,
+    CSVParser,
+)
+from torchvision.prototype.datasets.decoder import raw
+from torchvision.prototype.datasets.utils import (
+    Dataset,
+    DatasetConfig,
+    DatasetInfo,
+    HttpResource,
+    OnlineResource,
+    DatasetType,
+)
+from torchvision.prototype.datasets.utils._internal import INFINITE_BUFFER_SIZE, image_buffer_from_array
+
+
+class SEMEION(Dataset):
+    def _make_info(self) -> DatasetInfo:
+        return DatasetInfo(
+            "semeion",
+            type=DatasetType.RAW,
+            categories=10,
+            homepage="https://archive.ics.uci.edu/ml/datasets/Semeion+Handwritten+Digit",
+        )
+
+    def resources(self, config: DatasetConfig) -> List[OnlineResource]:
+        archive = HttpResource(
+            "http://archive.ics.uci.edu/ml/machine-learning-databases/semeion/semeion.data",
+            sha256="f43228ae3da5ea6a3c95069d53450b86166770e3b719dcc333182128fe08d4b1",
+        )
+        return [archive]
+
+    def _collate_and_decode_sample(
+        self,
+        data: Tuple[str, ...],
+        *,
+        decoder: Optional[Callable[[io.IOBase], torch.Tensor]],
+    ) -> Dict[str, Any]:
+        image_data = torch.tensor([float(pixel) for pixel in data[:256]], dtype=torch.uint8).reshape(16, 16)
+        label_data = [int(label) for label in data[256:] if label]
+
+        if decoder is raw:
+            image = image_data.unsqueeze(0)
+        else:
+            image_buffer = image_buffer_from_array(image_data.numpy())
+            image = decoder(image_buffer) if decoder else image_buffer  # type: ignore[assignment]
+
+        label = next((idx for idx, one_hot_label in enumerate(label_data) if one_hot_label))
+        category = self.info.categories[label]
+        return dict(image=image, label=label, category=category)
+
+    def _make_datapipe(
+        self,
+        resource_dps: List[IterDataPipe],
+        *,
+        config: DatasetConfig,
+        decoder: Optional[Callable[[io.IOBase], torch.Tensor]],
+    ) -> IterDataPipe[Dict[str, Any]]:
+        dp = resource_dps[0]
+        dp = CSVParser(dp, delimiter=" ")
+        dp = Shuffler(dp, buffer_size=INFINITE_BUFFER_SIZE)
+        dp = Mapper(dp, self._collate_and_decode_sample, fn_kwargs=dict(decoder=decoder))
+        return dp

torchvision/prototype/datasets/_builtin/semeion.py

@@ -52,7 +52,7 @@ def parse_args(argv=None):
 
 
 if __name__ == "__main__":
-    args = parse_args(["-f", "sbd"])
+    args = parse_args()
 
     try:
         main(*args.names, force=args.force)