File size: 1,003 Bytes
7cc2ec6 |
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 |
import torch
from transformers import AutoImageProcessor, Swinv2ForImageClassification
import onnx
import onnxruntime as ort
# Load the model and processor
image_processor = AutoImageProcessor.from_pretrained("haywoodsloan/ai-image-detector-deploy")
model = Swinv2ForImageClassification.from_pretrained("haywoodsloan/ai-image-detector-deploy")
# Set the model to evaluation mode
model.eval()
# Create dummy input for tracing
dummy_input = torch.randn(1, 3, 256, 256) # Batch size of 1, 3 color channels, 256x256 image
# Export the model to ONNX
onnx_model_path = "model.onnx"
# torch.onnx.export(
# model,
# dummy_input,
# onnx_model_path,
# input_names=["pixel_values"],
# output_names=["logits"],
# opset_version=11,
# dynamic_axes={
# "pixel_values": {0: "batch_size"},
# "logits": {0: "batch_size"}
# }
# )
# Verify the ONNX model
# onnx_model = onnx.load(onnx_model_path)
# onnx.checker.check_model(onnx_model)
print("The ONNX model is valid.") |