File size: 2,145 Bytes
5ee43e9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 | import argparse
import logging
import time
import torch
from transformers import AutoImageProcessor, ConvNextForImageClassification
from datasets import load_dataset
import torch_neuronx # ensures Neuron backend is available
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
def main():
parser = argparse.ArgumentParser(
description="ConvNeXt image-classification with torch.compile on Neuron"
)
parser.add_argument(
"--model",
type=str,
default="facebook/convnext-tiny-224",
help="ConvNeXT model name on Hugging Face Hub",
)
parser.add_argument("--batch-size", type=int, default=1, help="Batch size")
args = parser.parse_args()
torch.set_default_dtype(torch.float32)
torch.manual_seed(42)
# Load dataset and pick an image
dataset = load_dataset("huggingface/cats-image")
image = dataset["test"]["image"][0]
# Load processor and model
processor = AutoImageProcessor.from_pretrained(args.model)
model = ConvNextForImageClassification.from_pretrained(
args.model, torch_dtype=torch.float32, attn_implementation="eager"
)
model.eval()
# Preprocess image
inputs = processor(images=image, return_tensors="pt")
# Pre-run once to fix shapes before compilation
with torch.no_grad():
outputs = model(**inputs)
# Compile forward pass (allow graph breaks to avoid instruction-limit)
model.forward = torch.compile(model.forward, backend="neuron", fullgraph=False)
# Warmup
warmup_start = time.time()
with torch.no_grad():
_ = model(**inputs)
warmup_time = time.time() - warmup_start
# Actual run
run_start = time.time()
with torch.no_grad():
outputs = model(**inputs)
run_time = time.time() - run_start
# Predicted ImageNet class
predicted_class_idx = outputs.logits.argmax(-1).item()
predicted_label = model.config.id2label[predicted_class_idx]
logger.info("Warmup: %.2f s, Run: %.4f s", warmup_time, run_time)
logger.info("Predicted label: %s", predicted_label)
if __name__ == "__main__":
main() |