| """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()) | |