import io
import json
import os
from pathlib import Path
from typing import Callable, Dict, List, Optional, Tuple, Union
import numpy as np
import pandas as pd
import requests
import torch
import torch.utils.data as torchdata
import torchvision.transforms as transforms
from datasets import load_dataset
from ..models import nn as mnn
from ..models.utils import top_n_accuracy
from ..utils._download_data import url_is_reachable
from ..utils.const import TINY_IMAGENET_MEAN, TINY_IMAGENET_STD
from ._noniid_partition import non_iid_partition_with_dirichlet_distribution
from ._ops import CategoricalLabelToTensor, ImageArrayToTensor, ImageTensorScale
from ._register import register_fed_dataset
from .fed_dataset import FedVisionDataset, VisionDataset
if os.environ.get("HF_ENDPOINT", None) is not None and (not url_is_reachable(os.environ["HF_ENDPOINT"])):
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
elif os.environ.get("HF_ENDPOINT", None) is None and (not url_is_reachable("https://huggingface.co")):
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
__all__ = ["TinyImageNet"]
[docs]@register_fed_dataset()
class TinyImageNet(FedVisionDataset):
"""Tiny ImageNet dataset.
The Tiny ImageNet dataset is a subset of the ImageNet dataset. It consists of 200 classes, each with 500 training
images and 50 validation images and 50 test images. The images are downsampled to 64x64 pixels.
The original dataset [1]_ contains the test images while the hugingface dataset [3]_ does not contain the test images.
We use the hugingface dataset [3]_ for simplicity, and treat the validation set as the test set.
Parameters
----------
datadir : Union[pathlib.Path, str], optional
Directory to store data.
If ``None``, use default directory.
num_clients : int, default 100
Number of clients.
alpha : float, default 0.5
Concentration parameter for the Dirichlet distribution.
transform : Union[str, Callable], default "none"
Transform to apply to data. Conventions:
``"none"`` means no transform, using TensorDataset.
seed : int, default 0
Random seed for data partitioning.
**extra_config : dict, optional
Extra configurations.
References
----------
.. [1] http://cs231n.stanford.edu
.. [2] https://kaggle.com/competitions/tiny-imagenet
.. [3] https://huggingface.co/datasets/zh-plus/tiny-imagenet
"""
__name__ = "TinyImageNet"
def __init__(
self,
datadir: Optional[Union[Path, str]] = None,
num_clients: int = 100,
alpha: float = 0.5,
transform: Optional[Union[str, Callable]] = "none",
seed: int = 0,
) -> None:
self.num_clients = num_clients
self.alpha = alpha
super().__init__(datadir=datadir, transform=transform, seed=seed)
def _preload(self, datadir: Optional[Union[str, Path]] = None) -> None:
"""Preload the dataset.
Parameters
----------
datadir : Union[pathlib.Path, str], optional
Directory to store data.
If ``None``, use default directory.
Returns
-------
None
"""
if datadir is None:
ds = load_dataset("zh-plus/tiny-imagenet")
else:
ds = load_dataset("zh-plus/tiny-imagenet", data_dir=str(datadir))
self.datadir = Path(datadir or "~/.cache/huggingface/datasets/zh-plus___tiny-imagenet").expanduser().resolve()
self.datadir.mkdir(parents=True, exist_ok=True)
self._dataset_info = json.loads(list(self.datadir.rglob("dataset_info.json"))[0].read_text())
self.DEFAULT_TRAIN_CLIENTS_NUM = self.num_clients
self.DEFAULT_TEST_CLIENTS_NUM = self.num_clients
self.DEFAULT_BATCH_SIZE = 20
self.DEFAULT_TRAIN_FILE = [item["filename"] for item in ds.cache_files["train"]]
self.DEFAULT_TEST_FILE = [item["filename"] for item in ds.cache_files["valid"]]
self._IMGAE = "image"
self._LABEL = "label"
# load wnid to label mapping from imagenet website
timeout = 3
try:
self._wnid2label = pd.read_csv(
io.StringIO(requests.get("https://image-net.org/data/words.txt", timeout=timeout).text), sep="\t", header=None
)
except requests.exceptions.RequestException:
self._wnid2label = pd.DataFrame(columns=["wnid", "label"])
self._wnid2label.columns = ["wnid", "label"]
self._wnid2label["label"] = self._wnid2label["label"].apply(lambda x: str(x).split(",")[0])
self._wnid2label = self._wnid2label.set_index("wnid")["label"].to_dict()
# set criterion
self.criterion = torch.nn.CrossEntropyLoss()
# set transforms for creating dataset
if self.transform is None:
# set dynamic transform for train set
self.transform = transforms.Compose(
[
transforms.AutoAugment(
policy=transforms.AutoAugmentPolicy.IMAGENET,
),
ImageTensorScale(),
transforms.Normalize(TINY_IMAGENET_MEAN, TINY_IMAGENET_STD),
]
)
self.target_transform = transforms.Compose([CategoricalLabelToTensor()])
# load data
print("Loading data...")
self._train_data_dict = {
self._IMGAE: np.array(
[np.moveaxis(np.asarray(item[self._IMGAE].convert("RGB")), [0, 1], [1, 2]) for item in ds["train"]]
),
self._LABEL: np.array([item[self._LABEL] for item in ds["train"]]),
}
self._test_data_dict = {
self._IMGAE: np.array(
[np.moveaxis(np.asarray(item[self._IMGAE].convert("RGB")), [0, 1], [1, 2]) for item in ds["valid"]]
),
self._LABEL: np.array([item[self._LABEL] for item in ds["valid"]]),
}
self._n_class = 200
# distribute data into clients
print("Distributing data...")
self.indices = {}
self.indices["train"] = non_iid_partition_with_dirichlet_distribution(
label_list=self._train_data_dict[self._LABEL],
client_num=self.num_clients,
classes=self.n_class,
alpha=self.alpha,
)
self.indices["test"] = non_iid_partition_with_dirichlet_distribution(
label_list=self._test_data_dict[self._LABEL],
client_num=self.num_clients,
classes=self.n_class,
alpha=self.alpha,
)
[docs] def get_dataloader(
self,
train_bs: Optional[int] = None,
test_bs: Optional[int] = None,
client_idx: Optional[int] = None,
) -> Tuple[torchdata.DataLoader, torchdata.DataLoader]:
"""Get local dataloader at client `client_idx` or get the global dataloader.
Parameters
----------
train_bs : int, optional
Batch size for training dataloader.
If ``None``, use default batch size.
test_bs : int, optional
Batch size for testing dataloader.
If ``None``, use default batch size.
client_idx : int, optional
Index of the client to get dataloader.
If ``None``, get the dataloader containing all data.
Usually used for centralized training.
Returns
-------
train_dl : :class:`torch.utils.data.DataLoader`
Training dataloader.
test_dl : :class:`torch.utils.data.DataLoader`
Testing dataloader.
"""
if client_idx is None:
train_slice = slice(None)
test_slice = slice(None)
else:
train_slice = self.indices["train"][client_idx]
test_slice = self.indices["test"][client_idx]
# static transform
static_transform = transforms.Compose(
[
ImageArrayToTensor(),
ImageTensorScale(),
transforms.Normalize(TINY_IMAGENET_MEAN, TINY_IMAGENET_STD),
]
)
if self.transform == "none":
# apply only static transform
train_ds = torchdata.TensorDataset(
static_transform(self._train_data_dict[self._IMGAE][train_slice].copy()),
self.target_transform(self._train_data_dict[self._LABEL][train_slice].copy()),
)
else:
# use non-trivial dynamic transform
train_ds = VisionDataset(
images=torch.from_numpy(self._train_data_dict[self._IMGAE][train_slice].copy()).to(torch.uint8),
targets=self.target_transform(self._train_data_dict[self._LABEL][train_slice].copy()),
transform=self.transform,
)
train_dl = torchdata.DataLoader(
dataset=train_ds,
batch_size=train_bs or self.DEFAULT_BATCH_SIZE,
shuffle=True,
drop_last=False,
)
test_ds = torchdata.TensorDataset(
static_transform(self._test_data_dict[self._IMGAE][test_slice].copy()),
self.target_transform(self._test_data_dict[self._LABEL][test_slice].copy()),
)
test_dl = torchdata.DataLoader(
dataset=test_ds,
batch_size=test_bs or self.DEFAULT_BATCH_SIZE,
shuffle=False,
drop_last=False,
)
return train_dl, test_dl
[docs] def evaluate(self, probs: torch.Tensor, truths: torch.Tensor) -> Dict[str, float]:
"""Evaluation using predictions and ground truth.
Parameters
----------
probs : torch.Tensor
Predicted probabilities.
truths : torch.Tensor
Ground truth labels.
Returns
-------
Dict[str, float]
Evaluation results.
"""
return {
"acc": top_n_accuracy(probs, truths, 1),
"top3_acc": top_n_accuracy(probs, truths, 3),
"top5_acc": top_n_accuracy(probs, truths, 5),
"loss": self.criterion(probs, truths).item(),
"num_samples": probs.shape[0],
}
@property
def candidate_models(self) -> Dict[str, torch.nn.Module]:
"""A set of candidate models."""
return {
"resnet10": mnn.ResNet10(num_classes=self.n_class),
"resnet18": mnn.ResNet18(num_classes=self.n_class),
}
@property
def doi(self) -> List[str]:
"""DOI(s) related to the dataset."""
return ["10.1109/cvpr.2009.5206848"]
@property
def url(self) -> str:
"""URL for downloading the original dataset."""
return "http://cs231n.stanford.edu/tiny-imagenet-200.zip"
@property
def label_map(self) -> dict:
"""Label map for the dataset."""
return {
idx: self._wnid2label.get(label, label)
for idx, label in enumerate(self._dataset_info["features"]["label"]["names"])
}
[docs] def view_image(self, client_idx: int, image_idx: int) -> None:
"""View a single image.
Parameters
----------
client_idx : int
Index of the client on which the image is located.
image_idx : int
Index of the image in the client.
Returns
-------
None
"""
import matplotlib.pyplot as plt
if client_idx >= self.num_clients:
raise ValueError(f"client_idx must be less than {self.num_clients}, got {client_idx}")
total_num_images = len(self.indices["train"][client_idx]) + len(self.indices["test"][client_idx])
if image_idx >= total_num_images:
raise ValueError(f"image_idx must be less than {total_num_images}, got {image_idx}")
if image_idx < len(self.indices["train"][client_idx]):
image = self._train_data_dict[self._IMGAE][self.indices["train"][client_idx][image_idx]]
label = self._train_data_dict[self._LABEL][self.indices["train"][client_idx][image_idx]]
image_idx = self.indices["train"][client_idx][image_idx]
else:
image_idx -= len(self.indices["train"][client_idx])
image = self._test_data_dict[self._IMGAE][self.indices["test"][client_idx][image_idx]]
label = self._test_data_dict[self._LABEL][self.indices["test"][client_idx][image_idx]]
image_idx = self.indices["test"][client_idx][image_idx]
# image: channel first to channel last
image = image.transpose(1, 2, 0)
plt.imshow(image)
plt.title(f"image_idx: {image_idx}, label: {label} ({self.label_map[int(label)]}")
plt.show()
[docs] def random_grid_view(self, nrow: int, ncol: int, save_path: Optional[Union[str, Path]] = None) -> None:
"""Select randomly `nrow` x `ncol` images from the dataset
and plot them in a grid.
Parameters
----------
nrow : int
Number of rows in the grid.
ncol : int
Number of columns in the grid.
save_path : Union[str, Path], optional
Path to save the figure. If ``None``, do not save the figure.
Returns
-------
None
"""
import matplotlib.pyplot as plt
rng = np.random.default_rng()
fig, axes = plt.subplots(nrow, ncol, figsize=(ncol * 1, nrow * 1))
selected = []
for i in range(nrow):
for j in range(ncol):
while True:
client_idx = rng.integers(self.num_clients)
image_idx = rng.integers(len(self.indices["train"][client_idx]))
if (client_idx, image_idx) not in selected:
selected.append((client_idx, image_idx))
break
image = self._train_data_dict[self._IMGAE][self.indices["train"][client_idx][image_idx]]
axes[i, j].imshow(image.transpose(1, 2, 0))
axes[i, j].axis("off")
if save_path is not None:
fig.savefig(save_path, bbox_inches="tight", dpi=600)
plt.tight_layout()
plt.show()