wehub-resource-sync
121 lines
3.7 KiB
Python
121 lines
3.7 KiB
Python
import argparse
|
|
import io
|
|
import uuid
|
|
from typing import Any, Dict
|
|
|
|
import boto3
|
|
import numpy as np
|
|
import pandas as pd
|
|
import torch
|
|
from benchmark import Benchmark
|
|
from PIL import Image
|
|
from torchvision.models import ResNet50_Weights, resnet50
|
|
|
|
import ray
|
|
from ray.data import ActorPoolStrategy
|
|
|
|
BUCKET = "anyscale-imagenet"
|
|
# This Parquet file contains the keys of images in the 'anyscale-imagenet' bucket.
|
|
METADATA_PATH = "s3://anyscale-imagenet/metadata.parquet"
|
|
|
|
# Largest batch that can fit on a T4.
|
|
BATCH_SIZE = 900
|
|
|
|
WRITE_PATH = f"s3://ray-data-write-benchmark/{uuid.uuid4().hex}"
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument(
|
|
"--sf",
|
|
dest="scale_factor",
|
|
type=int,
|
|
default=1,
|
|
help=(
|
|
"The number of copies of ImageNet to read. Use this to simulate a larger "
|
|
"dataset."
|
|
),
|
|
)
|
|
return parser.parse_args()
|
|
|
|
|
|
def main(args: argparse.Namespace):
|
|
benchmark = Benchmark()
|
|
|
|
metadata = pd.read_parquet(METADATA_PATH)
|
|
# Repeat the metadata 'scale_factor' times to simulate a larger dataset.
|
|
metadata = pd.concat([metadata] * args.scale_factor, ignore_index=True)
|
|
|
|
def benchmark_fn():
|
|
weights = ResNet50_Weights.DEFAULT
|
|
model = resnet50(weights=weights)
|
|
model_ref = ray.put(model)
|
|
|
|
# Get the preprocessing transforms from the pre-trained weights.
|
|
transform = weights.transforms()
|
|
|
|
(
|
|
ray.data.from_pandas(metadata)
|
|
# TODO: There should be a way to specify "use as many actors as possible"
|
|
# with the now-recommended `concurrency` parameter.
|
|
.map(LoadImage, compute=ActorPoolStrategy(min_size=1))
|
|
# Preprocess the images using standard preprocessing
|
|
.map(ApplyTransform(transform))
|
|
.map_batches(
|
|
Predictor,
|
|
batch_size=BATCH_SIZE,
|
|
compute=ActorPoolStrategy(min_size=1),
|
|
num_gpus=1,
|
|
fn_constructor_kwargs={"model": model_ref, "device": "cuda"},
|
|
)
|
|
.write_parquet(WRITE_PATH)
|
|
)
|
|
|
|
benchmark.run_fn("main", benchmark_fn)
|
|
benchmark.write_result()
|
|
|
|
|
|
class LoadImage:
|
|
def __init__(self):
|
|
self._client = boto3.client("s3")
|
|
|
|
def __call__(self, row):
|
|
data = io.BytesIO()
|
|
self._client.download_fileobj(BUCKET, row["key"], data)
|
|
image = Image.open(data).convert("RGB")
|
|
return {"image": np.array(image)}
|
|
|
|
|
|
class ApplyTransform:
|
|
def __init__(self, transform):
|
|
self._transform = transform
|
|
|
|
def __call__(self, row: Dict[str, Any]) -> Dict[str, Any]:
|
|
# 'row["image"]' isn't writeable, and Torch only supports writeable tensors, so
|
|
# we need to maky a copy to prevent Torch from complaining.
|
|
tensor_batch = torch.as_tensor(np.copy(row["image"]), dtype=torch.float)
|
|
# (H, W, C) -> (C, H, W). This is required for the torchvision transform.
|
|
# https://pytorch.org/vision/main/models/generated/torchvision.models.resnet50.html#torchvision.models.ResNet50_Weights # noqa: E501
|
|
tensor_batch = tensor_batch.permute(2, 0, 1)
|
|
transformed_batch = self._transform(tensor_batch).numpy()
|
|
return {"image": transformed_batch}
|
|
|
|
|
|
class Predictor:
|
|
def __init__(self, model, device):
|
|
self._model = ray.get(model)
|
|
self._model.eval()
|
|
self._model.to(device)
|
|
|
|
self._device = device
|
|
|
|
def __call__(self, batch: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
|
|
with torch.inference_mode():
|
|
output = self._model(torch.as_tensor(batch["image"], device=self._device))
|
|
return {"predictions": output.cpu().numpy()}
|
|
|
|
|
|
if __name__ == "__main__":
|
|
args = parse_args()
|
|
main(args)
|