wehub-resource-sync
215 lines
7.6 KiB
Python
215 lines
7.6 KiB
Python
# Standard library imports
|
|
import logging
|
|
from typing import Dict
|
|
|
|
# Third-party imports
|
|
import torchvision
|
|
from torch.utils.data import IterableDataset
|
|
import pyarrow.fs
|
|
|
|
# Ray imports
|
|
import ray.train
|
|
from ray.data.datasource.partitioning import Partitioning
|
|
|
|
# Local imports
|
|
from constants import DatasetKey
|
|
from config import BenchmarkConfig
|
|
from image_classification.factory import (
|
|
ImageClassificationRayDataLoaderFactory,
|
|
ImageClassificationTorchDataLoaderFactory,
|
|
)
|
|
from image_classification.imagenet import get_transform
|
|
from s3_reader import AWS_REGION
|
|
from .imagenet import get_preprocess_map_fn
|
|
from .jpeg_iterable_dataset import S3JpegImageIterableDataset
|
|
from s3_jpeg_reader import S3JpegReader
|
|
from logger_utils import ContextLoggerAdapter
|
|
|
|
logger = ContextLoggerAdapter(logging.getLogger(__name__))
|
|
|
|
|
|
class ImageClassificationJpegRayDataLoaderFactory(
|
|
ImageClassificationRayDataLoaderFactory
|
|
):
|
|
"""Factory for creating Ray DataLoader for JPEG image classification.
|
|
|
|
Extends ImageClassificationRayDataLoaderFactory to provide:
|
|
1. S3 filesystem configuration with boto credentials
|
|
2. Ray dataset creation with partitioning by class
|
|
3. Resource allocation for concurrent validation
|
|
4. Image preprocessing with optional random transforms
|
|
"""
|
|
|
|
def __init__(self, benchmark_config: BenchmarkConfig, dataset_dirs: Dict[str, str]):
|
|
super().__init__(benchmark_config)
|
|
self._dataset_dirs = dataset_dirs
|
|
|
|
def get_s3fs_with_boto_creds(
|
|
self, connection_timeout: int = 60, request_timeout: int = 60
|
|
) -> pyarrow.fs.S3FileSystem:
|
|
"""Create S3 filesystem with boto credentials.
|
|
|
|
Args:
|
|
connection_timeout: Timeout for establishing connection in seconds
|
|
request_timeout: Timeout for requests in seconds
|
|
|
|
Returns:
|
|
Configured S3FileSystem instance with boto credentials
|
|
"""
|
|
import boto3
|
|
|
|
credentials = boto3.Session().get_credentials()
|
|
|
|
s3fs = pyarrow.fs.S3FileSystem(
|
|
access_key=credentials.access_key,
|
|
secret_key=credentials.secret_key,
|
|
session_token=credentials.token,
|
|
region=AWS_REGION,
|
|
connect_timeout=connection_timeout,
|
|
request_timeout=request_timeout,
|
|
)
|
|
return s3fs
|
|
|
|
def get_ray_datasets(self) -> Dict[str, ray.data.Dataset]:
|
|
"""Get Ray datasets for training and validation.
|
|
|
|
Creates training and validation datasets with:
|
|
1. Partitioning by class for efficient data loading
|
|
2. Image preprocessing with optional random transforms
|
|
3. Resource allocation for concurrent validation
|
|
4. Row limits based on benchmark configuration
|
|
|
|
Returns:
|
|
Dictionary containing:
|
|
- "train": Training dataset with random transforms
|
|
- "val": Validation dataset without transforms
|
|
"""
|
|
train_dir = self._dataset_dirs[DatasetKey.TRAIN]
|
|
# TODO: The validation dataset directory is not partitioned by class.
|
|
val_dir = train_dir
|
|
|
|
filesystem = (
|
|
self.get_s3fs_with_boto_creds() if train_dir.startswith("s3://") else None
|
|
)
|
|
|
|
# Create training dataset with class-based partitioning
|
|
train_partitioning = Partitioning(
|
|
"dir", base_dir=train_dir, field_names=["class"]
|
|
)
|
|
train_ds = ray.data.read_images(
|
|
train_dir,
|
|
mode="RGB",
|
|
include_paths=False,
|
|
partitioning=train_partitioning,
|
|
filesystem=filesystem,
|
|
).map(get_preprocess_map_fn(random_transforms=True))
|
|
|
|
if self.get_dataloader_config().limit_training_rows > 0:
|
|
train_ds = train_ds.limit(self.get_dataloader_config().limit_training_rows)
|
|
|
|
# Create validation dataset with same partitioning
|
|
val_partitioning = Partitioning("dir", base_dir=val_dir, field_names=["class"])
|
|
val_ds = ray.data.read_images(
|
|
val_dir,
|
|
mode="RGB",
|
|
include_paths=False,
|
|
partitioning=val_partitioning,
|
|
filesystem=filesystem,
|
|
).map(get_preprocess_map_fn(random_transforms=False))
|
|
|
|
if self.get_dataloader_config().limit_validation_rows > 0:
|
|
val_ds = val_ds.limit(self.get_dataloader_config().limit_validation_rows)
|
|
|
|
return {
|
|
DatasetKey.TRAIN: train_ds,
|
|
DatasetKey.VALID: val_ds,
|
|
}
|
|
|
|
|
|
class ImageClassificationJpegTorchDataLoaderFactory(
|
|
ImageClassificationTorchDataLoaderFactory, S3JpegReader
|
|
):
|
|
"""Factory for creating PyTorch DataLoaders for JPEG image classification.
|
|
|
|
Features:
|
|
- S3-based JPEG file reading with round-robin worker distribution
|
|
- Device transfer and error handling for data batches
|
|
- Row limits per worker for controlled processing
|
|
- Dataset caching for efficiency
|
|
"""
|
|
|
|
def __init__(self, benchmark_config: BenchmarkConfig, data_dirs: Dict[str, str]):
|
|
super().__init__(benchmark_config)
|
|
S3JpegReader.__init__(self) # Initialize S3JpegReader to set up _s3_client
|
|
self._data_dirs = data_dirs
|
|
self._cached_datasets = None
|
|
|
|
def get_iterable_datasets(self) -> Dict[str, IterableDataset]:
|
|
"""Get train and validation datasets with worker-specific configurations.
|
|
|
|
Returns:
|
|
Dictionary containing:
|
|
- "train": Training dataset with random transforms
|
|
- "val": Validation dataset without transforms
|
|
"""
|
|
if self._cached_datasets is not None:
|
|
return self._cached_datasets
|
|
|
|
if self._data_dirs[DatasetKey.TRAIN].startswith("s3://"):
|
|
return self._get_iterable_datasets_s3()
|
|
else:
|
|
return self._get_iterable_datasets_local()
|
|
|
|
def _get_iterable_datasets_local(self) -> Dict[str, IterableDataset]:
|
|
"""Get train and validation datasets from local filesystem."""
|
|
train_dir = self._data_dirs[DatasetKey.TRAIN]
|
|
val_dir = self._data_dirs[DatasetKey.VALID]
|
|
|
|
train_dataset = torchvision.datasets.ImageFolder(
|
|
root=train_dir,
|
|
transform=get_transform(to_torch_tensor=True, random_transforms=True),
|
|
)
|
|
|
|
val_dataset = torchvision.datasets.ImageFolder(
|
|
root=val_dir,
|
|
transform=get_transform(to_torch_tensor=True, random_transforms=False),
|
|
)
|
|
|
|
return {
|
|
DatasetKey.TRAIN: train_dataset,
|
|
DatasetKey.VALID: val_dataset,
|
|
}
|
|
|
|
def _get_iterable_datasets_s3(self) -> Dict[str, IterableDataset]:
|
|
"""Get train and validation datasets from S3."""
|
|
|
|
train_dir = self._data_dirs[DatasetKey.TRAIN]
|
|
|
|
# Get row limits for workers and total processing
|
|
(
|
|
limit_training_rows_per_worker,
|
|
limit_validation_rows_per_worker,
|
|
) = self._get_worker_row_limits()
|
|
|
|
# Get file URLs for training and validation
|
|
train_file_urls = val_file_urls = self._get_file_urls(train_dir)
|
|
train_ds = S3JpegImageIterableDataset(
|
|
file_urls=train_file_urls,
|
|
random_transforms=True,
|
|
limit_rows_per_worker=limit_training_rows_per_worker,
|
|
)
|
|
|
|
# TODO: IMAGENET_JPEG_SPLIT_S3_DIRS["val"] does not have the label
|
|
# partitioning like "train" does. So we use "train" for validation.
|
|
val_ds = S3JpegImageIterableDataset(
|
|
file_urls=val_file_urls,
|
|
random_transforms=False,
|
|
limit_rows_per_worker=limit_validation_rows_per_worker,
|
|
)
|
|
|
|
self._cached_datasets = {
|
|
DatasetKey.TRAIN: train_ds,
|
|
DatasetKey.VALID: val_ds,
|
|
}
|
|
return self._cached_datasets
|