chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,394 @@
|
||||
# Standard library imports
|
||||
import logging
|
||||
import time
|
||||
from typing import Dict, Tuple, Iterator, Generator, Optional, Union
|
||||
|
||||
# Third-party imports
|
||||
import torch
|
||||
import torchvision
|
||||
import pyarrow
|
||||
import ray
|
||||
import ray.train
|
||||
from ray.data.collate_fn import ArrowBatchCollateFn, CollateFn
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from ray.data.dataset import TorchDeviceType
|
||||
|
||||
# Local imports
|
||||
from benchmark_factory import BenchmarkFactory
|
||||
from config import BenchmarkConfig, DataloaderType, ImageClassificationConfig
|
||||
from dataloader_factory import BaseDataLoaderFactory
|
||||
from torch_dataloader_factory import TorchDataLoaderFactory
|
||||
from ray_dataloader_factory import RayDataLoaderFactory
|
||||
from logger_utils import ContextLoggerAdapter
|
||||
|
||||
logger = ContextLoggerAdapter(logging.getLogger(__name__))
|
||||
|
||||
|
||||
def mock_dataloader(
|
||||
num_batches: int = 64, batch_size: int = 32
|
||||
) -> Generator[Tuple[torch.Tensor, torch.Tensor], None, None]:
|
||||
"""Generate mock image and label tensors for testing.
|
||||
|
||||
Args:
|
||||
num_batches: Number of batches to generate
|
||||
batch_size: Number of samples per batch
|
||||
|
||||
Yields:
|
||||
Tuple of (image_tensor, label_tensor) for each batch
|
||||
"""
|
||||
device = ray.train.torch.get_device()
|
||||
|
||||
images = torch.randn(batch_size, 3, 224, 224).to(device)
|
||||
labels = torch.randint(0, 1000, (batch_size,)).to(device)
|
||||
|
||||
for _ in range(num_batches):
|
||||
yield images, labels
|
||||
|
||||
|
||||
class ImageClassificationTorchDataLoaderFactory(TorchDataLoaderFactory):
|
||||
"""Factory for creating PyTorch DataLoaders for image classification tasks.
|
||||
|
||||
Features:
|
||||
- Distributed file reading with round-robin worker distribution
|
||||
- Device transfer and error handling for data batches
|
||||
- Configurable row limits per worker for controlled processing
|
||||
- Performance monitoring and logging
|
||||
"""
|
||||
|
||||
def __init__(self, benchmark_config: BenchmarkConfig):
|
||||
super().__init__(benchmark_config)
|
||||
|
||||
def _calculate_rows_per_worker(
|
||||
self, total_rows: int, num_workers: int
|
||||
) -> Optional[int]:
|
||||
"""Calculate rows per worker for balanced data distribution.
|
||||
|
||||
Args:
|
||||
total_rows: Total rows to process across all workers (-1 for unlimited)
|
||||
num_workers: Total workers (Ray workers × Torch workers)
|
||||
|
||||
Returns:
|
||||
Rows per worker or None if no limit. Each worker gets at least 1 row.
|
||||
"""
|
||||
if total_rows < 0:
|
||||
return None
|
||||
|
||||
if num_workers == 0:
|
||||
return total_rows
|
||||
|
||||
return max(1, total_rows // num_workers)
|
||||
|
||||
def _get_worker_row_limits(self) -> Tuple[Optional[int], Optional[int]]:
|
||||
"""Calculate row limits per worker for training and validation.
|
||||
|
||||
Returns:
|
||||
Tuple of (training_rows_per_worker, validation_rows_per_worker)
|
||||
"""
|
||||
dataloader_config = self.get_dataloader_config()
|
||||
num_workers = max(1, dataloader_config.num_torch_workers)
|
||||
total_workers = self.benchmark_config.num_workers * num_workers
|
||||
|
||||
limit_training_rows_per_worker = self._calculate_rows_per_worker(
|
||||
self.get_dataloader_config().limit_training_rows, total_workers
|
||||
)
|
||||
|
||||
limit_validation_rows_per_worker = self._calculate_rows_per_worker(
|
||||
self.get_dataloader_config().limit_validation_rows, total_workers
|
||||
)
|
||||
|
||||
return limit_training_rows_per_worker, limit_validation_rows_per_worker
|
||||
|
||||
def create_batch_iterator(
|
||||
self, dataloader: torch.utils.data.DataLoader, device: torch.device
|
||||
) -> Iterator[Tuple[torch.Tensor, torch.Tensor]]:
|
||||
"""Create iterator with device transfer and error handling.
|
||||
|
||||
Args:
|
||||
dataloader: PyTorch DataLoader to iterate over
|
||||
device: Target device for tensor transfer
|
||||
|
||||
Returns:
|
||||
Iterator yielding (image_tensor, label_tensor) on target device
|
||||
"""
|
||||
worker_rank = ray.train.get_context().get_world_rank()
|
||||
logger.info(f"Worker {worker_rank}: Starting batch iteration")
|
||||
|
||||
try:
|
||||
last_batch_time = time.time()
|
||||
for batch_idx, batch in enumerate(dataloader):
|
||||
try:
|
||||
# Monitor batch processing delays
|
||||
current_time = time.time()
|
||||
time_since_last_batch = current_time - last_batch_time
|
||||
if time_since_last_batch > 10:
|
||||
logger.warning(
|
||||
f"Worker {worker_rank}: Long delay ({time_since_last_batch:.2f}s) "
|
||||
f"between batches {batch_idx-1} and {batch_idx}"
|
||||
)
|
||||
|
||||
# Process and transfer batch to device
|
||||
images, labels = batch
|
||||
logger.info(
|
||||
f"Worker {worker_rank}: Processing batch {batch_idx} (shape: {images.shape}, "
|
||||
f"time since last: {time_since_last_batch:.2f}s)"
|
||||
)
|
||||
|
||||
# Transfer tensors to target device
|
||||
transfer_start = time.time()
|
||||
dataloader_config = self.get_dataloader_config()
|
||||
images = images.to(
|
||||
device, non_blocking=dataloader_config.torch_non_blocking
|
||||
)
|
||||
labels = labels.to(
|
||||
device, non_blocking=dataloader_config.torch_non_blocking
|
||||
)
|
||||
transfer_time = time.time() - transfer_start
|
||||
|
||||
# Monitor device transfer performance
|
||||
if transfer_time > 5:
|
||||
logger.warning(
|
||||
f"Worker {worker_rank}: Slow device transfer ({transfer_time:.2f}s) "
|
||||
f"for batch {batch_idx}"
|
||||
)
|
||||
|
||||
logger.info(
|
||||
f"Worker {worker_rank}: Completed device transfer for batch {batch_idx} in "
|
||||
f"{transfer_time:.2f}s"
|
||||
)
|
||||
|
||||
last_batch_time = time.time()
|
||||
yield images, labels
|
||||
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Worker {worker_rank}: Error processing batch {batch_idx}: {str(e)}",
|
||||
exc_info=True,
|
||||
)
|
||||
raise
|
||||
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Worker {worker_rank}: Error in batch iterator: {str(e)}",
|
||||
exc_info=True,
|
||||
)
|
||||
raise
|
||||
|
||||
|
||||
class CustomArrowCollateFn(ArrowBatchCollateFn):
|
||||
"""Custom collate function for converting Arrow batches to PyTorch tensors."""
|
||||
|
||||
_DEFAULT_NUM_WORKERS = 4
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dtypes: Optional[Union["torch.dtype", Dict[str, "torch.dtype"]]] = None,
|
||||
device: Optional["TorchDeviceType"] = None,
|
||||
pin_memory: bool = False,
|
||||
num_workers: int = _DEFAULT_NUM_WORKERS,
|
||||
):
|
||||
"""Initialize the collate function.
|
||||
|
||||
Args:
|
||||
dtypes: Optional torch dtype(s) for the tensors
|
||||
device: Optional device to place tensors on
|
||||
pin_memory: Whether to pin the memory of the created tensors
|
||||
num_workers: Number of worker threads for parallel tensor conversion
|
||||
Defaults to `_DEFAULT_NUM_WORKERS`.
|
||||
"""
|
||||
import torch
|
||||
|
||||
self.dtypes = dtypes
|
||||
if isinstance(device, (str, int)):
|
||||
self.device = torch.device(device)
|
||||
else:
|
||||
self.device = device
|
||||
self.pin_memory = pin_memory
|
||||
self.num_workers = num_workers
|
||||
self._threadpool: Optional[ThreadPoolExecutor] = None
|
||||
|
||||
def __del__(self):
|
||||
"""Clean up threadpool on destruction."""
|
||||
if getattr(self, "_threadpool", None):
|
||||
self._threadpool.shutdown(wait=False)
|
||||
|
||||
def __call__(self, batch: "pyarrow.Table") -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Convert an Arrow batch to PyTorch tensors.
|
||||
|
||||
Args:
|
||||
batch: PyArrow Table to convert
|
||||
|
||||
Returns:
|
||||
Tuple of (image_tensor, label_tensor)
|
||||
"""
|
||||
from ray.data.util.torch_utils import (
|
||||
arrow_batch_to_tensors,
|
||||
)
|
||||
|
||||
if self.num_workers > 0 and self._threadpool is None:
|
||||
self._threadpool = ThreadPoolExecutor(max_workers=self.num_workers)
|
||||
|
||||
# For GPU transfer, we can skip the combining chunked arrays. This is because
|
||||
# we can convert the chunked arrays to corresponding numpy format and then to
|
||||
# Tensors and transfer the corresponding list of Tensors to GPU directly.
|
||||
# However, for CPU transfer, we need to combine the chunked arrays first
|
||||
# before converting to numpy format and then to Tensors.
|
||||
combine_chunks = self.device is not None and self.device.type == "cpu"
|
||||
tensors = arrow_batch_to_tensors(
|
||||
batch,
|
||||
dtypes=self.dtypes,
|
||||
combine_chunks=combine_chunks,
|
||||
pin_memory=self.pin_memory,
|
||||
threadpool=self._threadpool,
|
||||
)
|
||||
return tensors["image"], tensors["label"]
|
||||
|
||||
|
||||
class ImageClassificationRayDataLoaderFactory(RayDataLoaderFactory):
|
||||
"""Factory for creating Ray DataLoader for image classification tasks."""
|
||||
|
||||
def __init__(self, benchmark_config: BenchmarkConfig):
|
||||
super().__init__(benchmark_config)
|
||||
|
||||
def _get_collate_fn(self) -> Optional[CollateFn]:
|
||||
return CustomArrowCollateFn(
|
||||
device=ray.train.torch.get_device(),
|
||||
pin_memory=self.get_dataloader_config().ray_data_pin_memory,
|
||||
)
|
||||
|
||||
|
||||
class ImageClassificationMockDataLoaderFactory(BaseDataLoaderFactory):
|
||||
"""Factory for creating mock dataloaders for testing.
|
||||
|
||||
Provides mock implementations of training and validation dataloaders
|
||||
that generate random image and label tensors.
|
||||
"""
|
||||
|
||||
def get_train_dataloader(
|
||||
self,
|
||||
) -> Generator[Tuple[torch.Tensor, torch.Tensor], None, None]:
|
||||
"""Get mock training dataloader.
|
||||
|
||||
Returns:
|
||||
Generator yielding (image_tensor, label_tensor) batches
|
||||
"""
|
||||
dataloader_config = self.get_dataloader_config()
|
||||
return mock_dataloader(
|
||||
num_batches=1024, batch_size=dataloader_config.train_batch_size
|
||||
)
|
||||
|
||||
def get_val_dataloader(
|
||||
self,
|
||||
) -> Generator[Tuple[torch.Tensor, torch.Tensor], None, None]:
|
||||
"""Get mock validation dataloader.
|
||||
|
||||
Returns:
|
||||
Generator yielding (image_tensor, label_tensor) batches
|
||||
"""
|
||||
dataloader_config = self.get_dataloader_config()
|
||||
return mock_dataloader(
|
||||
num_batches=512, batch_size=dataloader_config.validation_batch_size
|
||||
)
|
||||
|
||||
|
||||
def get_imagenet_data_dirs(task_config: ImageClassificationConfig) -> Dict[str, str]:
|
||||
"""Returns a dict with the root imagenet dataset directories for train/val/test,
|
||||
corresponding to the data format and local/s3 dataset location."""
|
||||
from image_classification.imagenet import IMAGENET_LOCALFS_SPLIT_DIRS
|
||||
from image_classification.jpeg.imagenet import (
|
||||
IMAGENET_JPEG_SPLIT_S3_DIRS,
|
||||
)
|
||||
from image_classification.parquet.imagenet import (
|
||||
IMAGENET_PARQUET_SPLIT_S3_DIRS,
|
||||
IMAGENET_PARQUET_SPLIT_1T_S3_DIRS,
|
||||
)
|
||||
from image_classification.s3_url.imagenet import (
|
||||
IMAGENET_S3_URL_SPLIT_DIRS,
|
||||
)
|
||||
|
||||
data_format = task_config.image_classification_data_format
|
||||
|
||||
if task_config.image_classification_local_dataset:
|
||||
return IMAGENET_LOCALFS_SPLIT_DIRS
|
||||
|
||||
if data_format == ImageClassificationConfig.ImageFormat.JPEG:
|
||||
return IMAGENET_JPEG_SPLIT_S3_DIRS
|
||||
elif data_format == ImageClassificationConfig.ImageFormat.PARQUET:
|
||||
if task_config.image_classification_use_1t_dataset:
|
||||
return IMAGENET_PARQUET_SPLIT_1T_S3_DIRS
|
||||
return IMAGENET_PARQUET_SPLIT_S3_DIRS
|
||||
elif data_format == ImageClassificationConfig.ImageFormat.S3_URL:
|
||||
return IMAGENET_S3_URL_SPLIT_DIRS
|
||||
else:
|
||||
raise ValueError(f"Unknown data format: {data_format}")
|
||||
|
||||
|
||||
class ImageClassificationFactory(BenchmarkFactory):
|
||||
def get_dataloader_factory(self) -> BaseDataLoaderFactory:
|
||||
dataloader_type = self.benchmark_config.dataloader_type
|
||||
task_config = self.benchmark_config.task_config
|
||||
assert isinstance(task_config, ImageClassificationConfig)
|
||||
|
||||
data_dirs = get_imagenet_data_dirs(task_config)
|
||||
|
||||
data_format = task_config.image_classification_data_format
|
||||
|
||||
if dataloader_type == DataloaderType.MOCK:
|
||||
return ImageClassificationMockDataLoaderFactory(self.benchmark_config)
|
||||
|
||||
elif dataloader_type == DataloaderType.RAY_DATA:
|
||||
if data_format == ImageClassificationConfig.ImageFormat.JPEG:
|
||||
from image_classification.jpeg.factory import (
|
||||
ImageClassificationJpegRayDataLoaderFactory,
|
||||
)
|
||||
|
||||
return ImageClassificationJpegRayDataLoaderFactory(
|
||||
self.benchmark_config, data_dirs
|
||||
)
|
||||
elif data_format == ImageClassificationConfig.ImageFormat.PARQUET:
|
||||
from image_classification.parquet.factory import (
|
||||
ImageClassificationParquetRayDataLoaderFactory,
|
||||
)
|
||||
|
||||
return ImageClassificationParquetRayDataLoaderFactory(
|
||||
self.benchmark_config, data_dirs
|
||||
)
|
||||
elif data_format == ImageClassificationConfig.ImageFormat.S3_URL:
|
||||
# NOTE: This format downloads images via ray data expressions,
|
||||
# which is less efficient than native Ray Data S3 reading (JPEG format or Parquet format).
|
||||
# Use this primarily for testing the S3 URL download pattern.
|
||||
from image_classification.s3_url.factory import (
|
||||
ImageClassificationS3UrlRayDataLoaderFactory,
|
||||
)
|
||||
|
||||
return ImageClassificationS3UrlRayDataLoaderFactory(
|
||||
self.benchmark_config, data_dirs
|
||||
)
|
||||
|
||||
elif dataloader_type == DataloaderType.TORCH:
|
||||
if data_format == ImageClassificationConfig.ImageFormat.JPEG:
|
||||
from image_classification.jpeg.factory import (
|
||||
ImageClassificationJpegTorchDataLoaderFactory,
|
||||
)
|
||||
|
||||
return ImageClassificationJpegTorchDataLoaderFactory(
|
||||
self.benchmark_config, data_dirs
|
||||
)
|
||||
elif data_format == ImageClassificationConfig.ImageFormat.PARQUET:
|
||||
from image_classification.parquet.factory import (
|
||||
ImageClassificationParquetTorchDataLoaderFactory,
|
||||
)
|
||||
|
||||
return ImageClassificationParquetTorchDataLoaderFactory(
|
||||
self.benchmark_config, data_dirs
|
||||
)
|
||||
|
||||
raise ValueError(
|
||||
f"Invalid dataloader configuration: {dataloader_type}\n"
|
||||
f"{task_config}\n{self.benchmark_config.dataloader_config}"
|
||||
)
|
||||
|
||||
def get_model(self) -> torch.nn.Module:
|
||||
return torchvision.models.resnet50(weights=None)
|
||||
|
||||
def get_loss_fn(self) -> torch.nn.Module:
|
||||
return torch.nn.CrossEntropyLoss()
|
||||
File diff suppressed because it is too large
Load Diff
+53
@@ -0,0 +1,53 @@
|
||||
#!/bin/bash
|
||||
|
||||
set -e # Exit on any error
|
||||
|
||||
DATA_DIR="/mnt/local_storage/imagenet"
|
||||
ZIP_NAME="imagenet-64k.zip"
|
||||
ZIP_URL="s3://anyscale-imagenet/ILSVRC/Data/CLS-LOC/$ZIP_NAME"
|
||||
ZIP_PATH="$DATA_DIR/$ZIP_NAME"
|
||||
TRAIN_DIR="$DATA_DIR/train"
|
||||
|
||||
if [ ! -d "$TRAIN_DIR" ]; then
|
||||
echo "Downloading and extracting ImageNet subset to $DATA_DIR..."
|
||||
mkdir -p "$DATA_DIR"
|
||||
pushd "$DATA_DIR" || exit
|
||||
|
||||
echo "Fetching $ZIP_URL..."
|
||||
aws s3 cp "$ZIP_URL" "$ZIP_PATH"
|
||||
|
||||
echo "Unzipping..."
|
||||
unzip -q "$ZIP_NAME"
|
||||
rm "$ZIP_NAME"
|
||||
|
||||
popd || exit
|
||||
else
|
||||
echo "Dataset already exists at $TRAIN_DIR. Skipping download and unzip."
|
||||
fi
|
||||
|
||||
echo "Duplicating images in-place..."
|
||||
|
||||
python3 <<EOF
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
from tqdm import tqdm
|
||||
|
||||
dataset_dir = Path("$TRAIN_DIR")
|
||||
|
||||
for class_dir in tqdm(sorted(dataset_dir.iterdir()), desc="Processing classes"):
|
||||
if not class_dir.is_dir():
|
||||
continue
|
||||
|
||||
# Skip if already duplicated
|
||||
if any(class_dir.glob("*_copy1.JPEG")):
|
||||
print(f"Skipping {class_dir.name} (already duplicated)")
|
||||
continue
|
||||
|
||||
for img_path in class_dir.glob("*.JPEG"):
|
||||
for i in range(1, 8):
|
||||
copy_name = img_path.stem + f"_copy{i}" + img_path.suffix
|
||||
copy_path = class_dir / copy_name
|
||||
shutil.copy2(img_path, copy_path)
|
||||
EOF
|
||||
|
||||
echo "Image duplication complete."
|
||||
@@ -0,0 +1,214 @@
|
||||
# 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
|
||||
@@ -0,0 +1,57 @@
|
||||
import torch
|
||||
import numpy as np
|
||||
from typing import Dict, Union, Callable
|
||||
from PIL import Image
|
||||
|
||||
from constants import DatasetKey
|
||||
from image_classification.imagenet import (
|
||||
get_transform,
|
||||
IMAGENET_WNID_TO_ID,
|
||||
)
|
||||
|
||||
|
||||
IMAGENET_JPEG_SPLIT_S3_ROOT = "s3://anyscale-imagenet/ILSVRC/Data/CLS-LOC"
|
||||
IMAGENET_JPEG_SPLIT_S3_DIRS = {
|
||||
DatasetKey.TRAIN: f"{IMAGENET_JPEG_SPLIT_S3_ROOT}/train",
|
||||
DatasetKey.VALID: f"{IMAGENET_JPEG_SPLIT_S3_ROOT}/val",
|
||||
DatasetKey.TEST: f"{IMAGENET_JPEG_SPLIT_S3_ROOT}/test",
|
||||
}
|
||||
|
||||
|
||||
def get_preprocess_map_fn(
|
||||
random_transforms: bool = True,
|
||||
) -> Callable[[Dict[str, Union[np.ndarray, str]]], Dict[str, torch.Tensor]]:
|
||||
"""Get a map function that transforms a row to the format expected by the training loop.
|
||||
|
||||
Args:
|
||||
random_transforms: Whether to use random transforms for training
|
||||
|
||||
Returns:
|
||||
A function that takes a row dict with:
|
||||
- "image": numpy array in HWC format
|
||||
- "class": WNID string
|
||||
|
||||
The output is a dict with "image" and "label" keys.
|
||||
"""
|
||||
crop_resize_transform = get_transform(
|
||||
to_torch_tensor=True, random_transforms=random_transforms
|
||||
)
|
||||
|
||||
def map_fn(row: Dict[str, Union[np.ndarray, str]]) -> Dict[str, torch.Tensor]:
|
||||
"""Process a single row into the expected format.
|
||||
|
||||
Args:
|
||||
row: Dict containing "image" and "class" keys
|
||||
|
||||
Returns:
|
||||
Dict with "image" and "label" keys
|
||||
"""
|
||||
# Convert NumPy HWC image to PIL
|
||||
image_pil = Image.fromarray(row["image"])
|
||||
# Apply transform (includes ToTensor + Normalize)
|
||||
image = crop_resize_transform(image_pil)
|
||||
# Convert label
|
||||
label = IMAGENET_WNID_TO_ID[row["class"]]
|
||||
return {"image": image, "label": label}
|
||||
|
||||
return map_fn
|
||||
@@ -0,0 +1,266 @@
|
||||
# Standard library imports
|
||||
import io
|
||||
import logging
|
||||
import time
|
||||
from typing import Iterator, List, Optional, Tuple, Callable
|
||||
|
||||
# Third-party imports
|
||||
import numpy as np
|
||||
from PIL import Image as PILImage
|
||||
import torch
|
||||
from torch.utils.data import IterableDataset
|
||||
|
||||
# Ray imports
|
||||
import ray
|
||||
import ray.train
|
||||
|
||||
# Local imports
|
||||
from s3_jpeg_reader import S3JpegReader
|
||||
from .imagenet import get_preprocess_map_fn
|
||||
from logger_utils import ContextLoggerAdapter
|
||||
|
||||
logger = ContextLoggerAdapter(logging.getLogger(__name__))
|
||||
|
||||
|
||||
class S3JpegImageIterableDataset(S3JpegReader, IterableDataset):
|
||||
"""An iterable dataset that loads images from S3-stored JPEG files.
|
||||
|
||||
Features:
|
||||
- Direct image fetching from S3
|
||||
- Optional random transforms for training
|
||||
- Row limits per worker for controlled processing
|
||||
- Progress logging and performance metrics
|
||||
"""
|
||||
|
||||
# Constants
|
||||
LOG_FREQUENCY = 1000 # Log progress every 1000 rows
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
file_urls: List[str],
|
||||
random_transforms: bool = True,
|
||||
limit_rows_per_worker: Optional[int] = None,
|
||||
):
|
||||
"""Initialize the dataset.
|
||||
|
||||
Args:
|
||||
file_urls: List of S3 URLs to load
|
||||
random_transforms: Whether to use random transforms for training
|
||||
limit_rows_per_worker: Maximum number of rows to process per worker
|
||||
"""
|
||||
super().__init__()
|
||||
self.file_urls = file_urls
|
||||
self.limit_rows_per_worker = limit_rows_per_worker
|
||||
self.random_transforms = random_transforms
|
||||
|
||||
worker_rank = ray.train.get_context().get_world_rank()
|
||||
logger.info(
|
||||
f"Worker {worker_rank}: Initialized with {len(file_urls)} files"
|
||||
f"{f' (limit: {limit_rows_per_worker} rows)' if limit_rows_per_worker else ''}"
|
||||
)
|
||||
|
||||
def _get_worker_info(self) -> Tuple[int, int]:
|
||||
"""Get current worker information.
|
||||
|
||||
Returns:
|
||||
Tuple of (worker_id, num_workers)
|
||||
"""
|
||||
worker_info = torch.utils.data.get_worker_info()
|
||||
worker_id = worker_info.id if worker_info else 0
|
||||
num_workers = worker_info.num_workers if worker_info else 1
|
||||
return worker_id, num_workers
|
||||
|
||||
def _has_reached_row_limit(self, rows_processed: int) -> bool:
|
||||
"""Check if we've reached the row limit per worker.
|
||||
|
||||
Args:
|
||||
rows_processed: Number of rows processed so far
|
||||
|
||||
Returns:
|
||||
True if we've reached the limit, False otherwise
|
||||
"""
|
||||
return (
|
||||
self.limit_rows_per_worker is not None
|
||||
and rows_processed >= self.limit_rows_per_worker
|
||||
)
|
||||
|
||||
def _log_progress(
|
||||
self, worker_id: int, rows_processed: int, last_log_time: float
|
||||
) -> float:
|
||||
"""Log processing progress and return updated last_log_time.
|
||||
|
||||
Args:
|
||||
worker_id: ID of the current worker
|
||||
rows_processed: Number of rows processed so far
|
||||
last_log_time: Time of last progress log
|
||||
|
||||
Returns:
|
||||
Updated last_log_time
|
||||
"""
|
||||
if rows_processed % self.LOG_FREQUENCY == 0:
|
||||
current_time = time.time()
|
||||
elapsed_time = current_time - last_log_time
|
||||
rows_per_second = (
|
||||
self.LOG_FREQUENCY / elapsed_time if elapsed_time > 0 else 0
|
||||
)
|
||||
logger.info(
|
||||
f"Worker {worker_id}: Processed {rows_processed} rows "
|
||||
f"({rows_per_second:.2f} rows/sec)"
|
||||
)
|
||||
return current_time
|
||||
return last_log_time
|
||||
|
||||
def _fetch_image(self, file_url: str) -> Tuple[str, Optional[PILImage.Image]]:
|
||||
"""Fetch a single image from S3.
|
||||
|
||||
Args:
|
||||
file_url: S3 URL to fetch (e.g., "s3://bucket/path/to/image.jpg")
|
||||
|
||||
Returns:
|
||||
Tuple of (file_url, PIL Image) where PIL Image is None if fetch failed
|
||||
"""
|
||||
worker_id, _ = self._get_worker_info()
|
||||
try:
|
||||
bucket = file_url.replace("s3://", "").split("/")[0]
|
||||
key = "/".join(file_url.replace("s3://", "").split("/")[1:])
|
||||
|
||||
response = self.s3_client.get_object(Bucket=bucket, Key=key)
|
||||
image_data = response["Body"].read()
|
||||
image = PILImage.open(io.BytesIO(image_data))
|
||||
return file_url, image
|
||||
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Worker {worker_id}: Error fetching image from {file_url}: {str(e)}",
|
||||
exc_info=True,
|
||||
)
|
||||
return file_url, None
|
||||
|
||||
def _process_image(
|
||||
self,
|
||||
image: PILImage.Image,
|
||||
file_url: str,
|
||||
preprocess_fn: Callable,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Process a single image and convert to tensors.
|
||||
|
||||
Args:
|
||||
image: PIL Image to process
|
||||
file_url: URL of the image file
|
||||
preprocess_fn: Preprocessing function to apply
|
||||
|
||||
Returns:
|
||||
Tuple of (image_tensor, label_tensor)
|
||||
"""
|
||||
try:
|
||||
# Convert to RGB and numpy array
|
||||
if image.mode != "RGB":
|
||||
image = image.convert("RGB")
|
||||
image_array = np.array(image, dtype=np.uint8)
|
||||
|
||||
# Ensure HWC format (Height x Width x Channels)
|
||||
if len(image_array.shape) == 2: # Grayscale
|
||||
image_array = np.stack([image_array] * 3, axis=-1)
|
||||
elif len(image_array.shape) == 3 and image_array.shape[0] == 3: # CHW
|
||||
image_array = np.transpose(image_array, (1, 2, 0))
|
||||
|
||||
wnid = file_url.split("/")[-2] # Extract WNID from path
|
||||
|
||||
processed = preprocess_fn({"image": image_array, "class": wnid})
|
||||
|
||||
image = torch.as_tensor(processed["image"], dtype=torch.float32)
|
||||
label = torch.as_tensor(processed["label"], dtype=torch.int64)
|
||||
|
||||
return image, label
|
||||
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Error processing {file_url}: {str(e)}",
|
||||
exc_info=True,
|
||||
)
|
||||
raise
|
||||
|
||||
def _process_files(
|
||||
self, files_to_read: List[str], preprocess_fn: Callable, worker_id: int
|
||||
) -> Iterator[Tuple[torch.Tensor, torch.Tensor]]:
|
||||
"""Process multiple files and yield processed rows.
|
||||
|
||||
Args:
|
||||
files_to_read: List of file URLs to process
|
||||
preprocess_fn: Preprocessing function to apply
|
||||
worker_id: ID of the current worker
|
||||
|
||||
Yields:
|
||||
Tuple of (image_tensor, label_tensor)
|
||||
"""
|
||||
rows_processed = 0
|
||||
last_log_time = time.time()
|
||||
total_start_time = time.time()
|
||||
|
||||
for file_url in files_to_read:
|
||||
if self._has_reached_row_limit(rows_processed):
|
||||
logger.info(f"Worker {worker_id}: Reached row limit: {rows_processed}")
|
||||
break
|
||||
|
||||
file_url, image = self._fetch_image(file_url)
|
||||
if image is None:
|
||||
continue
|
||||
|
||||
try:
|
||||
image, label = self._process_image(
|
||||
image,
|
||||
file_url,
|
||||
preprocess_fn,
|
||||
)
|
||||
rows_processed += 1
|
||||
last_log_time = self._log_progress(
|
||||
worker_id, rows_processed, last_log_time
|
||||
)
|
||||
yield image, label
|
||||
except Exception:
|
||||
continue
|
||||
|
||||
# Log final statistics
|
||||
total_time = time.time() - total_start_time
|
||||
logger.info(
|
||||
f"Worker {worker_id}: Finished: {rows_processed} rows in {total_time:.2f}s "
|
||||
f"({rows_processed/total_time:.2f} rows/sec)"
|
||||
)
|
||||
|
||||
def __iter__(self) -> Iterator[Tuple[torch.Tensor, torch.Tensor]]:
|
||||
"""Iterate through the dataset and yield (image, label) tensors.
|
||||
|
||||
Yields:
|
||||
Tuple[torch.Tensor, torch.Tensor]: (image, label) tensors
|
||||
|
||||
Raises:
|
||||
Exception: If there's a fatal error during iteration
|
||||
"""
|
||||
try:
|
||||
# Get worker info for file distribution
|
||||
worker_id, num_workers = self._get_worker_info()
|
||||
logger.info(f"Worker {worker_id}/{num_workers}: Starting")
|
||||
|
||||
# Distribute files among workers
|
||||
files_to_read = (
|
||||
self.file_urls
|
||||
if num_workers == 1
|
||||
else self.file_urls[worker_id::num_workers]
|
||||
)
|
||||
|
||||
logger.info(f"Worker {worker_id}: Processing {len(files_to_read)} files")
|
||||
|
||||
# Initialize preprocessing function
|
||||
preprocess_fn = get_preprocess_map_fn(
|
||||
random_transforms=self.random_transforms
|
||||
)
|
||||
|
||||
# Process files and yield results
|
||||
yield from self._process_files(files_to_read, preprocess_fn, worker_id)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Worker {worker_id}: Fatal error: {str(e)}",
|
||||
exc_info=True,
|
||||
)
|
||||
raise
|
||||
@@ -0,0 +1,139 @@
|
||||
# Standard library imports
|
||||
import logging
|
||||
from typing import Dict, Optional
|
||||
|
||||
# Third-party imports
|
||||
from torch.utils.data import IterableDataset
|
||||
import ray
|
||||
import ray.data
|
||||
import ray.train
|
||||
|
||||
# Local imports
|
||||
from constants import DatasetKey
|
||||
from config import BenchmarkConfig
|
||||
from image_classification.factory import (
|
||||
ImageClassificationRayDataLoaderFactory,
|
||||
ImageClassificationTorchDataLoaderFactory,
|
||||
)
|
||||
from .imagenet import get_preprocess_map_fn
|
||||
from .parquet_iterable_dataset import S3ParquetImageIterableDataset
|
||||
from s3_parquet_reader import S3ParquetReader
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ImageClassificationParquetRayDataLoaderFactory(
|
||||
ImageClassificationRayDataLoaderFactory
|
||||
):
|
||||
"""Factory for creating Ray DataLoader for Parquet image classification.
|
||||
|
||||
Features:
|
||||
- Parquet file reading with column selection
|
||||
- Image decoding and preprocessing
|
||||
- Resource allocation for concurrent validation
|
||||
- Row limits based on benchmark configuration
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, benchmark_config: BenchmarkConfig, data_dirs: Dict[str, str]
|
||||
) -> None:
|
||||
super().__init__(benchmark_config)
|
||||
self._data_dirs = data_dirs
|
||||
|
||||
def get_ray_datasets(self) -> Dict[str, ray.data.Dataset]:
|
||||
"""Get Ray datasets for training and validation.
|
||||
|
||||
Returns:
|
||||
Dictionary containing:
|
||||
- "train": Training dataset with random transforms
|
||||
- "val": Validation dataset without transforms
|
||||
"""
|
||||
# Create training dataset with image decoding and transforms
|
||||
train_ds = ray.data.read_parquet(
|
||||
self._data_dirs[DatasetKey.TRAIN],
|
||||
columns=["image", "label"],
|
||||
).map(get_preprocess_map_fn(decode_image=True, 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 without random transforms
|
||||
val_ds = ray.data.read_parquet(
|
||||
self._data_dirs[DatasetKey.TRAIN],
|
||||
columns=["image", "label"],
|
||||
).map(get_preprocess_map_fn(decode_image=True, 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 ImageClassificationParquetTorchDataLoaderFactory(
|
||||
ImageClassificationTorchDataLoaderFactory, S3ParquetReader
|
||||
):
|
||||
"""Factory for creating PyTorch DataLoaders for Parquet image classification.
|
||||
|
||||
Features:
|
||||
- Parquet file reading with row count-based distribution
|
||||
- Worker-based file distribution for balanced workloads
|
||||
- Row limits per worker for controlled processing
|
||||
- Dataset instance caching for efficiency
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, benchmark_config: BenchmarkConfig, data_dirs: Dict[str, str]
|
||||
) -> None:
|
||||
"""Initialize factory with benchmark configuration.
|
||||
|
||||
Args:
|
||||
benchmark_config: Configuration for benchmark parameters
|
||||
"""
|
||||
super().__init__(benchmark_config)
|
||||
S3ParquetReader.__init__(
|
||||
self
|
||||
) # Initialize S3ParquetReader to set up _s3_client
|
||||
self.train_url = data_dirs[DatasetKey.TRAIN]
|
||||
self._cached_datasets: Optional[Dict[str, IterableDataset]] = 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
|
||||
|
||||
# Get row limits for workers and total processing
|
||||
(
|
||||
limit_training_rows_per_worker,
|
||||
limit_validation_rows_per_worker,
|
||||
) = self._get_worker_row_limits()
|
||||
|
||||
# Create training dataset
|
||||
train_file_urls = self._get_file_urls(self.train_url)
|
||||
train_ds = S3ParquetImageIterableDataset(
|
||||
file_urls=train_file_urls,
|
||||
random_transforms=True,
|
||||
limit_rows_per_worker=limit_training_rows_per_worker,
|
||||
)
|
||||
|
||||
# Create validation dataset
|
||||
val_file_urls = train_file_urls
|
||||
val_ds = S3ParquetImageIterableDataset(
|
||||
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
|
||||
@@ -0,0 +1,69 @@
|
||||
import io
|
||||
import numpy as np
|
||||
from typing import Dict, Union, Callable
|
||||
from PIL import Image
|
||||
from torchvision.transforms.functional import pil_to_tensor
|
||||
|
||||
from constants import DatasetKey
|
||||
from image_classification.imagenet import (
|
||||
get_transform,
|
||||
IMAGENET_WNID_TO_ID,
|
||||
)
|
||||
|
||||
IMAGENET_PARQUET_SPLIT_S3_ROOT = (
|
||||
"s3://ray-benchmark-data-internal-us-west-2/imagenet/parquet_split"
|
||||
)
|
||||
IMAGENET_PARQUET_SPLIT_S3_DIRS = {
|
||||
DatasetKey.TRAIN: f"{IMAGENET_PARQUET_SPLIT_S3_ROOT}/train",
|
||||
DatasetKey.VALID: f"{IMAGENET_PARQUET_SPLIT_S3_ROOT}/val",
|
||||
DatasetKey.TEST: f"{IMAGENET_PARQUET_SPLIT_S3_ROOT}/test",
|
||||
}
|
||||
|
||||
# Much larger parquet dataset used to sustain Ray Data backpressure in the
|
||||
# slow-consumer ingest benchmarks.
|
||||
IMAGENET_PARQUET_SPLIT_1T_S3_ROOT = (
|
||||
"s3://ray-benchmark-data-internal-us-west-2/imagenet/parquet_split_1t"
|
||||
)
|
||||
IMAGENET_PARQUET_SPLIT_1T_S3_DIRS = {
|
||||
DatasetKey.TRAIN: f"{IMAGENET_PARQUET_SPLIT_1T_S3_ROOT}/train",
|
||||
DatasetKey.VALID: f"{IMAGENET_PARQUET_SPLIT_1T_S3_ROOT}/val",
|
||||
DatasetKey.TEST: f"{IMAGENET_PARQUET_SPLIT_1T_S3_ROOT}/test",
|
||||
}
|
||||
|
||||
|
||||
def get_preprocess_map_fn(
|
||||
decode_image: bool = True, random_transforms: bool = True
|
||||
) -> Callable[[Dict[str, Union[bytes, str]]], Dict[str, Union[np.ndarray, int]]]:
|
||||
"""Get a map function that transforms a row of the dataset to the format
|
||||
expected by the training loop.
|
||||
|
||||
Args:
|
||||
decode_image: Whether to decode the image bytes into a tensor
|
||||
random_transforms: Whether to use random transforms for training
|
||||
|
||||
Returns:
|
||||
A function that takes a row dict and returns a processed dict.
|
||||
Input row dict should have:
|
||||
- "image": bytes or tensor in CHW format
|
||||
- "label": WNID string
|
||||
|
||||
Output dict has:
|
||||
- "image": np.array of the transformed, normalized image
|
||||
- "label": An integer index of the WNID
|
||||
"""
|
||||
crop_resize_transform = get_transform(
|
||||
to_torch_tensor=False, random_transforms=random_transforms
|
||||
)
|
||||
|
||||
def map_fn(row: Dict[str, Union[bytes, str]]) -> Dict[str, Union[np.ndarray, int]]:
|
||||
assert "image" in row and "label" in row, row.keys()
|
||||
|
||||
if decode_image:
|
||||
row["image"] = pil_to_tensor(Image.open(io.BytesIO(row["image"]))) / 255.0
|
||||
|
||||
row["image"] = np.array(crop_resize_transform(row["image"]))
|
||||
row["label"] = IMAGENET_WNID_TO_ID[row["label"]]
|
||||
|
||||
return {"image": row["image"], "label": row["label"]}
|
||||
|
||||
return map_fn
|
||||
+273
@@ -0,0 +1,273 @@
|
||||
# Standard library imports
|
||||
from typing import List, Tuple, Optional, Iterator, Callable
|
||||
import logging
|
||||
import io
|
||||
import time
|
||||
|
||||
# Third-party imports
|
||||
import pandas as pd
|
||||
import pyarrow.parquet as pq
|
||||
import torch
|
||||
from torch.utils.data import IterableDataset
|
||||
|
||||
# Ray imports
|
||||
import ray
|
||||
import ray.train
|
||||
|
||||
# Local imports
|
||||
from s3_parquet_reader import S3ParquetReader
|
||||
from .imagenet import get_preprocess_map_fn
|
||||
from logger_utils import ContextLoggerAdapter
|
||||
|
||||
logger = ContextLoggerAdapter(logging.getLogger(__name__))
|
||||
|
||||
|
||||
# TODO Look into https://github.com/webdataset/webdataset for more canonical way to do data
|
||||
# distribution between Ray Train and Torch Dataloader workers.
|
||||
|
||||
|
||||
class S3ParquetImageIterableDataset(S3ParquetReader, IterableDataset):
|
||||
"""An iterable dataset that loads images from S3-stored Parquet files.
|
||||
|
||||
This dataset:
|
||||
1. Reads Parquet files from S3 one row group at a time
|
||||
2. Processes images with optional random transforms
|
||||
3. Yields (image, label) tensors
|
||||
4. Supports row limits per worker for controlled data processing
|
||||
"""
|
||||
|
||||
LOG_FREQUENCY = 1000 # Log progress every 1000 rows
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
file_urls: List[str],
|
||||
random_transforms: bool = True,
|
||||
limit_rows_per_worker: Optional[int] = None,
|
||||
):
|
||||
"""Initialize the dataset.
|
||||
|
||||
Args:
|
||||
file_urls: List of S3 URLs to load
|
||||
random_transforms: Whether to use random transforms for training
|
||||
limit_rows_per_worker: Maximum number of rows to process per worker (None for all rows)
|
||||
"""
|
||||
super().__init__()
|
||||
self.file_urls = file_urls
|
||||
self.limit_rows_per_worker = limit_rows_per_worker
|
||||
self.random_transforms = random_transforms
|
||||
|
||||
worker_rank = ray.train.get_context().get_world_rank()
|
||||
logger.info(
|
||||
f"Worker {worker_rank}: Initialized with {len(file_urls)} files"
|
||||
f"{f' (limit: {limit_rows_per_worker} rows)' if limit_rows_per_worker else ''}"
|
||||
)
|
||||
|
||||
def _get_worker_info(self) -> Tuple[int, int]:
|
||||
"""Get current worker information.
|
||||
|
||||
Returns:
|
||||
Tuple of (worker_id, num_workers)
|
||||
"""
|
||||
worker_info = torch.utils.data.get_worker_info()
|
||||
worker_id = worker_info.id if worker_info else 0
|
||||
num_workers = worker_info.num_workers if worker_info else 1
|
||||
return worker_id, num_workers
|
||||
|
||||
def _has_reached_row_limit(self, rows_processed: int) -> bool:
|
||||
"""Check if we've reached the row limit per worker.
|
||||
|
||||
Args:
|
||||
rows_processed: Number of rows processed so far
|
||||
|
||||
Returns:
|
||||
True if we've reached the limit, False otherwise
|
||||
"""
|
||||
return (
|
||||
self.limit_rows_per_worker is not None
|
||||
and rows_processed >= self.limit_rows_per_worker
|
||||
)
|
||||
|
||||
def _log_progress(
|
||||
self, worker_id: int, rows_processed: int, last_log_time: float
|
||||
) -> float:
|
||||
"""Log processing progress and return updated last_log_time.
|
||||
|
||||
Args:
|
||||
worker_id: ID of the current worker
|
||||
rows_processed: Number of rows processed so far
|
||||
last_log_time: Time of last progress log
|
||||
|
||||
Returns:
|
||||
Updated last_log_time
|
||||
"""
|
||||
if rows_processed % self.LOG_FREQUENCY == 0:
|
||||
current_time = time.time()
|
||||
elapsed_time = current_time - last_log_time
|
||||
rows_per_second = (
|
||||
self.LOG_FREQUENCY / elapsed_time if elapsed_time > 0 else 0
|
||||
)
|
||||
logger.info(
|
||||
f"Worker {worker_id}: Processed {rows_processed} rows "
|
||||
f"({rows_per_second:.2f} rows/sec)"
|
||||
)
|
||||
return current_time
|
||||
return last_log_time
|
||||
|
||||
def _read_parquet_file(self, file_url: str) -> Iterator[pd.DataFrame]:
|
||||
"""Read a Parquet file from S3 one row group at a time.
|
||||
|
||||
This method:
|
||||
1. Fetches the Parquet file from S3
|
||||
2. Reads it row group by row group
|
||||
3. Converts each row group to a pandas DataFrame
|
||||
|
||||
Args:
|
||||
file_url: S3 URL of the Parquet file
|
||||
|
||||
Yields:
|
||||
DataFrame containing one row group at a time
|
||||
|
||||
Raises:
|
||||
Exception: If there's an error reading the file
|
||||
"""
|
||||
try:
|
||||
start_time = time.time()
|
||||
worker_id, _ = self._get_worker_info()
|
||||
logger.info(f"Worker {worker_id}: Reading Parquet file: {file_url}")
|
||||
|
||||
# Get parquet file metadata
|
||||
bucket, key = self._parse_s3_url(file_url)
|
||||
response = self.s3_client.get_object(Bucket=bucket, Key=key)
|
||||
parquet_file = pq.ParquetFile(io.BytesIO(response["Body"].read()))
|
||||
num_row_groups = parquet_file.num_row_groups
|
||||
|
||||
logger.info(
|
||||
f"Worker {worker_id}: Found {num_row_groups} row groups in {file_url}"
|
||||
)
|
||||
|
||||
for row_group in range(num_row_groups):
|
||||
# Read row group and convert to pandas
|
||||
table = parquet_file.read_row_group(row_group)
|
||||
df = table.to_pandas()
|
||||
yield df
|
||||
|
||||
total_time = time.time() - start_time
|
||||
logger.info(
|
||||
f"Worker {worker_id}: Completed reading {file_url} in {total_time:.2f}s"
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
worker_id, _ = self._get_worker_info()
|
||||
logger.error(
|
||||
f"Worker {worker_id}: Error reading file {file_url}: {str(e)}",
|
||||
exc_info=True,
|
||||
)
|
||||
raise
|
||||
|
||||
def _process_file(
|
||||
self,
|
||||
file_url: str,
|
||||
preprocess_fn: Callable,
|
||||
) -> Iterator[Tuple[torch.Tensor, torch.Tensor]]:
|
||||
"""Process a single file and yield processed rows.
|
||||
|
||||
Args:
|
||||
file_url: URL of the file to process
|
||||
preprocess_fn: Preprocessing function to apply
|
||||
|
||||
Yields:
|
||||
Tuple of (image_tensor, label_tensor)
|
||||
"""
|
||||
for df in self._read_parquet_file(file_url):
|
||||
for _, row in df.iterrows():
|
||||
try:
|
||||
# Process row and convert to tensors
|
||||
processed = preprocess_fn(row)
|
||||
image = torch.as_tensor(processed["image"], dtype=torch.float32)
|
||||
label = torch.as_tensor(processed["label"], dtype=torch.int64)
|
||||
yield image, label
|
||||
except Exception:
|
||||
continue
|
||||
|
||||
def _process_files(
|
||||
self, files_to_read: List[str], preprocess_fn: Callable, worker_id: int
|
||||
) -> Iterator[Tuple[torch.Tensor, torch.Tensor]]:
|
||||
"""Process multiple files and yield processed rows.
|
||||
|
||||
Args:
|
||||
files_to_read: List of file URLs to process
|
||||
preprocess_fn: Preprocessing function to apply
|
||||
worker_id: ID of the current worker
|
||||
|
||||
Yields:
|
||||
Tuple of (image_tensor, label_tensor)
|
||||
"""
|
||||
rows_processed = 0
|
||||
last_log_time = time.time()
|
||||
total_start_time = time.time()
|
||||
|
||||
for file_url in files_to_read:
|
||||
if self._has_reached_row_limit(rows_processed):
|
||||
logger.info(f"Worker {worker_id}: Reached row limit: {rows_processed}")
|
||||
break
|
||||
|
||||
for image, label in self._process_file(file_url, preprocess_fn):
|
||||
if self._has_reached_row_limit(rows_processed):
|
||||
break
|
||||
|
||||
rows_processed += 1
|
||||
last_log_time = self._log_progress(
|
||||
worker_id, rows_processed, last_log_time
|
||||
)
|
||||
yield image, label
|
||||
|
||||
# Log final statistics
|
||||
total_time = time.time() - total_start_time
|
||||
logger.info(
|
||||
f"Worker {worker_id}: Finished: {rows_processed} rows in {total_time:.2f}s "
|
||||
f"({rows_processed/total_time:.2f} rows/sec)"
|
||||
)
|
||||
|
||||
def __iter__(self) -> Iterator[Tuple[torch.Tensor, torch.Tensor]]:
|
||||
"""Main iteration method that processes files and yields (image, label) tensors.
|
||||
|
||||
This method:
|
||||
1. Distributes files among workers
|
||||
2. Processes rows with image transforms
|
||||
3. Converts to tensors
|
||||
4. Respects row limits per worker
|
||||
|
||||
Yields:
|
||||
Tuple of (image_tensor, label_tensor)
|
||||
|
||||
Raises:
|
||||
Exception: If there's a fatal error during processing
|
||||
"""
|
||||
try:
|
||||
# Get worker info for file distribution
|
||||
worker_id, num_workers = self._get_worker_info()
|
||||
logger.info(f"Worker {worker_id}/{num_workers}: Starting")
|
||||
|
||||
# Initialize preprocessing function
|
||||
preprocess_fn = get_preprocess_map_fn(
|
||||
decode_image=True, random_transforms=self.random_transforms
|
||||
)
|
||||
|
||||
# Distribute files among workers
|
||||
files_to_read = (
|
||||
self.file_urls
|
||||
if num_workers == 1
|
||||
else self.file_urls[worker_id::num_workers]
|
||||
)
|
||||
|
||||
logger.info(f"Worker {worker_id}: Processing {len(files_to_read)} files")
|
||||
|
||||
# Process files and yield results
|
||||
yield from self._process_files(files_to_read, preprocess_fn, worker_id)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Worker {worker_id}: Fatal error: {str(e)}",
|
||||
exc_info=True,
|
||||
)
|
||||
raise
|
||||
@@ -0,0 +1,77 @@
|
||||
# Standard library imports
|
||||
import logging
|
||||
from typing import Dict
|
||||
|
||||
# Third-party imports
|
||||
import ray.data
|
||||
|
||||
# Local imports
|
||||
from constants import DatasetKey
|
||||
from config import BenchmarkConfig
|
||||
from image_classification.factory import ImageClassificationRayDataLoaderFactory
|
||||
from .imagenet import (
|
||||
create_s3_url_dataset,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ImageClassificationS3UrlRayDataLoaderFactory(
|
||||
ImageClassificationRayDataLoaderFactory
|
||||
):
|
||||
"""Factory for creating Ray DataLoader that downloads images from S3 URLs.
|
||||
|
||||
This factory:
|
||||
1. Lists JPEG files from S3 using boto3
|
||||
2. Creates a Ray dataset from the file records
|
||||
3. Uses map_batches to download and process images from S3
|
||||
|
||||
This approach separates file listing from image downloading, which can be
|
||||
more efficient for certain workloads as it allows parallel downloads during
|
||||
the map_batches execution on CPU workers.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, benchmark_config: BenchmarkConfig, data_dirs: Dict[str, str]
|
||||
) -> None:
|
||||
super().__init__(benchmark_config)
|
||||
self._data_dirs = data_dirs
|
||||
|
||||
def get_ray_datasets(self) -> Dict[str, ray.data.Dataset]:
|
||||
"""Get Ray datasets for training and validation.
|
||||
|
||||
Returns:
|
||||
Dictionary containing:
|
||||
- "train": Training dataset with random transforms
|
||||
- "val": Validation dataset without transforms
|
||||
"""
|
||||
dataloader_config = self.get_dataloader_config()
|
||||
|
||||
# Create training dataset
|
||||
train_limit = (
|
||||
dataloader_config.limit_training_rows
|
||||
if dataloader_config.limit_training_rows > 0
|
||||
else None
|
||||
)
|
||||
train_ds = create_s3_url_dataset(
|
||||
data_dir=self._data_dirs[DatasetKey.TRAIN],
|
||||
random_transforms=True,
|
||||
limit_rows=train_limit,
|
||||
)
|
||||
|
||||
# Create validation dataset
|
||||
val_limit = (
|
||||
dataloader_config.limit_validation_rows
|
||||
if dataloader_config.limit_validation_rows > 0
|
||||
else None
|
||||
)
|
||||
val_ds = create_s3_url_dataset(
|
||||
data_dir=self._data_dirs[DatasetKey.TRAIN],
|
||||
random_transforms=False,
|
||||
limit_rows=val_limit,
|
||||
)
|
||||
|
||||
return {
|
||||
DatasetKey.TRAIN: train_ds,
|
||||
DatasetKey.VALID: val_ds,
|
||||
}
|
||||
@@ -0,0 +1,251 @@
|
||||
"""ImageNet dataset loading via S3 URL download with Ray Data expressions.
|
||||
|
||||
This module provides dataset loading that:
|
||||
1. Lists JPEG files from S3 using boto3 (parallelized via Ray tasks)
|
||||
2. Creates a Ray dataset from the file records
|
||||
3. Uses Ray Data expressions (alpha) to download image bytes efficiently
|
||||
4. Uses map_batches to decode and process images
|
||||
|
||||
This approach leverages Ray Data's expressions API for optimized parallel I/O,
|
||||
separating the download step from image processing for better throughput.
|
||||
"""
|
||||
|
||||
import io
|
||||
import logging
|
||||
from functools import lru_cache
|
||||
from typing import Callable, Dict, List, Optional, Tuple
|
||||
|
||||
import boto3
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
from torchvision.transforms.functional import pil_to_tensor
|
||||
|
||||
import ray.data
|
||||
from ray.data.expressions import download
|
||||
|
||||
from constants import DatasetKey
|
||||
from image_classification.imagenet import (
|
||||
get_transform,
|
||||
IMAGENET_WNID_TO_ID,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# S3 configuration for ImageNet JPEG data
|
||||
AWS_REGION = "us-west-2"
|
||||
S3_ROOT = "s3://anyscale-imagenet/ILSVRC/Data/CLS-LOC"
|
||||
IMAGENET_S3_URL_SPLIT_DIRS = {
|
||||
DatasetKey.TRAIN: f"{S3_ROOT}/train",
|
||||
DatasetKey.VALID: f"{S3_ROOT}/val",
|
||||
DatasetKey.TEST: f"{S3_ROOT}/test",
|
||||
}
|
||||
|
||||
|
||||
def _get_class_labels(bucket: str, prefix: str) -> List[str]:
|
||||
"""Get all class label directories from S3.
|
||||
|
||||
Args:
|
||||
bucket: S3 bucket name
|
||||
prefix: S3 prefix path
|
||||
|
||||
Returns:
|
||||
List of class label directory names
|
||||
"""
|
||||
from typing import Set
|
||||
|
||||
# Ensure prefix ends with /
|
||||
if prefix and not prefix.endswith("/"):
|
||||
prefix += "/"
|
||||
|
||||
# List directories using delimiter
|
||||
s3_client = boto3.client("s3", region_name=AWS_REGION)
|
||||
paginator = s3_client.get_paginator("list_objects_v2")
|
||||
|
||||
# Use delimiter to get "directory" level
|
||||
labels: Set[str] = set()
|
||||
for page in paginator.paginate(Bucket=bucket, Prefix=prefix, Delimiter="/"):
|
||||
# CommonPrefixes contains the "directories"
|
||||
for common_prefix in page.get("CommonPrefixes", []):
|
||||
prefix_path = common_prefix["Prefix"]
|
||||
# Extract the directory name
|
||||
label = prefix_path.rstrip("/").split("/")[-1]
|
||||
labels.add(label)
|
||||
|
||||
return sorted(labels)
|
||||
|
||||
|
||||
@ray.remote
|
||||
def _list_files_for_label(
|
||||
bucket: str, prefix: str, label: str
|
||||
) -> List[Tuple[str, str]]:
|
||||
"""Ray task to list all image files for a specific label.
|
||||
|
||||
Args:
|
||||
bucket: S3 bucket name
|
||||
prefix: S3 prefix (parent directory)
|
||||
label: Class label (subdirectory name)
|
||||
|
||||
Returns:
|
||||
List of tuples with (file_path, class_name)
|
||||
"""
|
||||
s3_client = boto3.client("s3", region_name=AWS_REGION)
|
||||
paginator = s3_client.get_paginator("list_objects_v2")
|
||||
|
||||
# Construct the full prefix for this label
|
||||
label_prefix = f"{prefix}/{label}/" if prefix else f"{label}/"
|
||||
|
||||
file_records = []
|
||||
for page in paginator.paginate(Bucket=bucket, Prefix=label_prefix):
|
||||
for obj in page.get("Contents", []):
|
||||
key = obj["Key"]
|
||||
if key.lower().endswith((".jpg", ".jpeg")):
|
||||
file_path = f"s3://{bucket}/{key}"
|
||||
file_records.append((file_path, label))
|
||||
|
||||
return file_records
|
||||
|
||||
|
||||
@lru_cache(maxsize=8)
|
||||
def _list_s3_image_files_cached(data_dir: str) -> Tuple[Tuple[str, str], ...]:
|
||||
"""Cached implementation of S3 file listing using Ray tasks for parallelism.
|
||||
|
||||
Returns a tuple of tuples for hashability (required by lru_cache).
|
||||
"""
|
||||
logger.info(f"Listing JPEG files from {data_dir}...")
|
||||
|
||||
# Parse S3 URL: s3://bucket/prefix
|
||||
s3_path = data_dir
|
||||
if s3_path.startswith("s3://"):
|
||||
s3_path = s3_path[5:]
|
||||
parts = s3_path.split("/", 1)
|
||||
bucket = parts[0]
|
||||
prefix = parts[1].rstrip("/") if len(parts) > 1 else ""
|
||||
|
||||
# Get all class labels
|
||||
labels = _get_class_labels(bucket, prefix)
|
||||
logger.info(
|
||||
f"Found {len(labels)} class labels, launching Ray tasks for parallel listing..."
|
||||
)
|
||||
|
||||
# Launch Ray tasks for each label
|
||||
futures = [_list_files_for_label.remote(bucket, prefix, label) for label in labels]
|
||||
|
||||
# Wait for all tasks to complete and aggregate results
|
||||
results = ray.get(futures)
|
||||
|
||||
# Flatten the list of lists
|
||||
file_records = []
|
||||
for records in results:
|
||||
file_records.extend(records)
|
||||
|
||||
logger.info(f"Listed and cached {len(file_records)} JPEG files")
|
||||
return tuple(file_records)
|
||||
|
||||
|
||||
def list_s3_image_files(data_dir: str) -> List[Dict[str, str]]:
|
||||
"""List JPEG files from S3 with class labels extracted from path.
|
||||
|
||||
Results are cached to avoid repeated S3 listings.
|
||||
|
||||
Args:
|
||||
data_dir: S3 path to list files from (e.g., "s3://bucket/prefix")
|
||||
|
||||
Returns:
|
||||
List of dicts with "path" (S3 URL) and "class" (WNID) keys
|
||||
"""
|
||||
cached_records = _list_s3_image_files_cached(data_dir)
|
||||
return [{"path": path, "class": cls} for path, cls in cached_records]
|
||||
|
||||
|
||||
def get_process_batch_fn(
|
||||
random_transforms: bool = True,
|
||||
label_to_id_map: Optional[Dict[str, int]] = None,
|
||||
) -> Callable[[Dict[str, np.ndarray]], Dict[str, np.ndarray]]:
|
||||
"""Get a map_batches function that processes pre-downloaded image bytes.
|
||||
|
||||
This function expects image bytes to already be downloaded (via Ray Data
|
||||
expressions) and handles decoding and transformations.
|
||||
|
||||
Args:
|
||||
random_transforms: Whether to use random transforms for training
|
||||
label_to_id_map: Mapping from WNID strings to integer IDs
|
||||
|
||||
Returns:
|
||||
A function suitable for use with dataset.map_batches()
|
||||
"""
|
||||
if label_to_id_map is None:
|
||||
label_to_id_map = IMAGENET_WNID_TO_ID
|
||||
|
||||
transform = get_transform(
|
||||
to_torch_tensor=False, random_transforms=random_transforms
|
||||
)
|
||||
|
||||
def process_batch(
|
||||
batch: Dict[str, np.ndarray],
|
||||
) -> Dict[str, np.ndarray]:
|
||||
"""Process pre-downloaded image bytes.
|
||||
|
||||
Args:
|
||||
batch: Dict with "bytes" (image data) and "class" arrays
|
||||
|
||||
Returns:
|
||||
Dict with "image" (numpy array) and "label" (int) arrays
|
||||
"""
|
||||
processed_images = []
|
||||
labels = []
|
||||
|
||||
image_bytes_list = list(batch["bytes"])
|
||||
classes = list(batch["class"])
|
||||
|
||||
for data, wnid in zip(image_bytes_list, classes):
|
||||
# Decode and transform image
|
||||
image_pil = Image.open(io.BytesIO(data)).convert("RGB")
|
||||
image_tensor = pil_to_tensor(image_pil) / 255.0
|
||||
processed_image = np.array(transform(image_tensor))
|
||||
processed_images.append(processed_image)
|
||||
|
||||
# Convert label
|
||||
labels.append(label_to_id_map[wnid])
|
||||
|
||||
return {
|
||||
"image": np.stack(processed_images),
|
||||
"label": np.array(labels),
|
||||
}
|
||||
|
||||
return process_batch
|
||||
|
||||
|
||||
def create_s3_url_dataset(
|
||||
data_dir: str,
|
||||
random_transforms: bool = True,
|
||||
limit_rows: Optional[int] = None,
|
||||
) -> ray.data.Dataset:
|
||||
"""Create a Ray dataset that downloads images from S3 URLs.
|
||||
|
||||
Uses Ray Data expressions (alpha) for efficient parallel downloads,
|
||||
then map_batches for image decoding and transformations.
|
||||
|
||||
Args:
|
||||
data_dir: S3 path to the image directory
|
||||
random_transforms: Whether to use random transforms
|
||||
limit_rows: Optional row limit
|
||||
|
||||
Returns:
|
||||
Ray dataset with "image" and "label" columns
|
||||
"""
|
||||
file_records = list_s3_image_files(data_dir)
|
||||
|
||||
ds = ray.data.from_items(file_records)
|
||||
|
||||
if limit_rows is not None and limit_rows > 0:
|
||||
ds = ds.limit(limit_rows)
|
||||
|
||||
# Download image bytes using Ray Data expressions (alpha)
|
||||
# This enables optimized parallel I/O managed by Ray Data
|
||||
ds = ds.with_column("bytes", download("path"))
|
||||
|
||||
# Process downloaded bytes (decode and transform)
|
||||
process_fn = get_process_batch_fn(random_transforms=random_transforms)
|
||||
ds = ds.map_batches(process_fn)
|
||||
|
||||
return ds
|
||||
Reference in New Issue
Block a user