Create utilities to load Image Classification models & datasets from HF

examples/src/armory/examples/image_classification/food101.py

@@ -20,6 +20,7 @@
 import armory.engine
 import armory.evaluation
 import armory.export.image_classification
+import armory.helpers.huggingface
 import armory.metric
 import armory.metrics.compute
 import armory.metrics.perturbation
@@ -73,70 +74,20 @@ def parse_cli_args():
 
 def load_model():
     """Load model from HuggingFace"""
-    hf_model = armory.track.track_params(
-        transformers.AutoModelForImageClassification.from_pretrained
-    )(pretrained_model_name_or_path="nateraw/food")
-
-    armory_model = armory.model.image_classification.ImageClassifier(
-        name="ViT-finetuned-food101",
-        model=hf_model,
-        accessor=armory.data.Images.as_torch(scale=normalized_scale),
-    )
-
-    art_classifier = armory.track.track_init_params(
-        art.estimators.classification.PyTorchClassifier
-    )(
-        armory_model,
-        loss=torch.nn.CrossEntropyLoss(),
-        optimizer=torch.optim.Adam(armory_model.parameters(), lr=0.003),
-        input_shape=(3, 224, 224),
-        channels_first=True,
-        nb_classes=101,
-        clip_values=(-1.0, 1.0),
-    )
-
-    return armory_model, art_classifier
-
-
-def transform(processor, sample):
-    """Use the HF image processor and convert from BW To RGB"""
-    sample["image"] = processor([img.convert("RGB") for img in sample["image"]])[
-        "pixel_values"
-    ]
-    return sample
+    return armory.helpers.huggingface.load_image_classification_model("nateraw/food")
 
 
 def load_huggingface_dataset(batch_size: int, shuffle: bool):
     """Load food-101 dataset from HuggingFace"""
-    import functools
-
-    import datasets
-
-    hf_dataset = datasets.load_dataset("food101", split="validation")
-    assert isinstance(hf_dataset, datasets.Dataset)
-
-    labels = hf_dataset.features["label"].names
-
     hf_processor = transformers.AutoImageProcessor.from_pretrained("nateraw/food")
-    hf_dataset.set_transform(functools.partial(transform, hf_processor))
-
-    dataloader = armory.dataset.ImageClassificationDataLoader(
-        hf_dataset,
-        dim=armory.data.ImageDimensions.CHW,
-        scale=normalized_scale,
-        image_key="image",
-        label_key="label",
+    return armory.helpers.huggingface.load_image_classification_dataset(
+        "food101",
+        split="validation",
+        processor=hf_processor,
         batch_size=batch_size,
         shuffle=shuffle,
     )
 
-    evaluation_dataset = armory.evaluation.Dataset(
-        name="food-101",
-        dataloader=dataloader,
-    )
-
-    return evaluation_dataset, labels
-
 
 def load_torchvision_dataset(
     batch_size: int, shuffle: bool, sysconfig: armory.evaluation.SysConfig
@@ -177,12 +128,13 @@ def load_torchvision_dataset(
         shuffle=shuffle,
     )
 
-    evaluation_dataset = armory.evaluation.Dataset(
+    evaluation_dataset = armory.evaluation.ImageClassificationDataset(
         name="food-101",
         dataloader=dataloader,
+        labels=labels,
     )
 
-    return evaluation_dataset, labels
+    return evaluation_dataset
 
 
 def create_pgd_attack(classifier: art.estimators.classification.PyTorchClassifier):
@@ -290,7 +242,7 @@ def main(
 
     sysconfig = armory.evaluation.SysConfig()
     model, art_classifier = load_model()
-    dataset, _ = (
+    dataset = (
         load_huggingface_dataset(batch_size, shuffle)
         if dataset_src == "huggingface"
         else load_torchvision_dataset(batch_size, shuffle, sysconfig)

library/src/armory/evaluation.py

@@ -1,4 +1,5 @@
 """Armory Experiment Configuration Classes"""
+
 from dataclasses import dataclass, field
 import os
 from pathlib import Path
@@ -9,6 +10,7 @@
     Mapping,
     Optional,
     Protocol,
+    Sequence,
     runtime_checkable,
 )
 
@@ -32,6 +34,12 @@ class Dataset:
     """Data loader for evaluation data"""
 
 
+@dataclass
+class ImageClassificationDataset(Dataset):
+    labels: Sequence[str]
+    """List of class label names"""
+
+
 @runtime_checkable
 class ModelProtocol(Protocol):
     """Model being evaluated"""

library/src/armory/helpers/huggingface.py

@@ -0,0 +1,145 @@
+"""
+Armory helper utilities to assist with use of HuggingFace datasets and models
+with Armory
+"""
+
+from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple
+
+from art.estimators.classification.pytorch import PyTorchClassifier
+from transformers import AutoImageProcessor
+
+from armory.data import DataType, ImageDimensions, Images, Scale
+from armory.evaluation import ImageClassificationDataset
+from armory.model.image_classification.image_classifier import ImageClassifier
+from armory.track import track_init_params, track_params
+
+if TYPE_CHECKING:
+    import torch.nn
+
+Sample = Dict[str, Any]
+
+
+def _create_scale_from_image_processor(processor: AutoImageProcessor) -> Scale:
+    do_normalize = getattr(processor, "do_normalize", False)
+    image_mean = getattr(processor, "image_mean", None)
+    image_std = getattr(processor, "image_std", None)
+    scale = Scale(
+        dtype=DataType.FLOAT,
+        max=1.0,
+        mean=image_mean if do_normalize else None,
+        std=image_std if do_normalize else None,
+    )
+    return scale
+
+
+def _create_clip_values_from_scale(scale: Scale) -> Tuple[float, float]:
+    if not scale.is_normalized or not scale.mean or not scale.std:
+        return (0, scale.max)
+    else:
+        min_val = 0.0 - max(scale.mean) / min(scale.std)
+        max_val = scale.max - min(scale.mean) / min(scale.std)
+        return (min_val, max_val)
+
+
+def _transform(processor: AutoImageProcessor, sample: Sample) -> Sample:
+    """Use the HF image processor and convert from BW To RGB"""
+    sample["image"] = processor([img.convert("RGB") for img in sample["image"]])[
+        "pixel_values"
+    ]
+    return sample
+
+
+@track_params
+def load_image_classification_dataset(
+    name: str,
+    split: str,
+    processor: AutoImageProcessor,
+    dim: ImageDimensions = ImageDimensions.CHW,
+    image_key: str = "image",
+    label_key: str = "label",
+    **kwargs,
+) -> ImageClassificationDataset:
+    import functools
+
+    import datasets
+
+    from armory.dataset import ImageClassificationDataLoader
+    from armory.evaluation import ImageClassificationDataset
+
+    hf_dataset = datasets.load_dataset(name, split=split)
+    assert isinstance(hf_dataset, datasets.Dataset)
+
+    labels = hf_dataset.features[label_key].names
+
+    hf_dataset.set_transform(functools.partial(_transform, processor))
+
+    scale = _create_scale_from_image_processor(processor)
+
+    dataloader = ImageClassificationDataLoader(
+        hf_dataset,
+        dim=dim,
+        scale=scale,
+        image_key=image_key,
+        label_key=label_key,
+        **kwargs,
+    )
+
+    evaluation_dataset = ImageClassificationDataset(
+        name=name,
+        dataloader=dataloader,
+        labels=labels,
+    )
+
+    return evaluation_dataset
+
+
+@track_params
+def load_image_classification_model(
+    name: str,
+    loss: Optional["torch.nn.modules.loss._Loss"] = None,
+    dim: ImageDimensions = ImageDimensions.CHW,
+    num_channels: int = 3,
+) -> Tuple[ImageClassifier, PyTorchClassifier]:
+    from art.estimators.classification import PyTorchClassifier
+    from transformers import AutoImageProcessor, AutoModelForImageClassification
+
+    from armory.model.image_classification import ImageClassifier
+
+    hf_model = AutoModelForImageClassification.from_pretrained(name)
+
+    hf_processor = AutoImageProcessor.from_pretrained(name)
+    if hf_processor is None:
+        raise RuntimeError(f"No image processor found for pretrained model, {name}")
+
+    scale = _create_scale_from_image_processor(hf_processor)
+    accessor = Images.as_torch(dim=dim, scale=scale)
+
+    armory_model = ImageClassifier(
+        name=name,
+        model=hf_model,
+        accessor=accessor,
+    )
+
+    if loss is None:
+        import torch.nn
+
+        loss = torch.nn.CrossEntropyLoss()
+
+    channels_first = dim == ImageDimensions.CHW
+    input_shape = (
+        (num_channels, hf_processor.size["height"], hf_processor.size["width"])
+        if channels_first
+        else (hf_processor.size["height"], hf_processor.size["width"], num_channels)
+    )
+    clip_values = _create_clip_values_from_scale(scale)
+
+    art_classifier = track_init_params(PyTorchClassifier)(
+        armory_model,
+        loss=loss,
+        input_shape=input_shape,
+        channels_first=channels_first,
+        nb_classes=hf_model.num_labels,
+        clip_values=clip_values,
+    )
+
+    return armory_model, art_classifier