File size: 984 Bytes
c855122 986a3a9 c855122 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 | """DINOv3 ViT-L/16 features on trn2 with vit-encoder-neuron-kernels.
NEURON_LOGICAL_NC_CONFIG=2 python dinov3_features.py
"""
import types
import torch
from kernels import get_kernel
from transformers import AutoModel
k = get_kernel("aws-neuron/vit-encoder-neuron-kernels", version=1, trust_remote_code=True)
model = AutoModel.from_pretrained("facebook/dinov3-vitl16-pretrain-lvd1689m", dtype=torch.float32).eval()
ref_model = AutoModel.from_pretrained("facebook/dinov3-vitl16-pretrain-lvd1689m", dtype=torch.float32).eval()
model = model.to("neuron")
model.forward = types.MethodType(k.NeuronDINOv3ViTModel.forward, model)
model = torch.compile(model, backend="neuron")
torch.manual_seed(0)
pixel_values = torch.randn(4, 3, 224, 224)
with torch.no_grad():
out = model(pixel_values.to("neuron")).last_hidden_state.cpu()
ref = ref_model(pixel_values).last_hidden_state
print("features", tuple(out.shape), "rel_l2 vs fp32 CPU:", ((out - ref).norm() / ref.norm()).item())
|