一个 DeiT 图像分类模型。由论文作者在 ImageNet-1k 上训练。
import torch
import torch_npu
from torch_npu.contrib import transfer_to_npu
from PIL import Image
from openmind import is_torch_npu_available
import timm
import argparse
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument(
"--model_name_or_path",
type=str,
help="Path to model",
default=None,
)
args = parser.parse_args()
return args
if __name__ == "__main__":
args = parse_args()
model_path = args.model_name_or_path
img_path = model_path + '/img/beignets-task-guide.png'
img = Image.open(img_path)
if is_torch_npu_available():
device = "npu:0"
else:
device = "cpu"
model_name = 'deit_small_patch16_224.fb_in1k'
checkpoint_path=model_path + '/pytorch_model.bin'
model = timm.create_model(model_name, pretrained=False, checkpoint_path=checkpoint_path).to(device)
model = model.eval()
# get model specific transforms (normalization, resize)
data_config = timm.data.resolve_model_data_config(model)
transforms = timm.data.create_transform(**data_config, is_training=False)
output = model(transforms(img).unsqueeze(0).to(device)) # unsqueeze single image into batch of 1
img.close()
top5_probabilities, top5_class_indices = torch.topk(output.softmax(dim=1) * 100, k=5)
print(top5_probabilities)
print(top5_class_indices)在 timm 的 model results 中探索此模型的数据集和运行时指标。
@InProceedings{pmlr-v139-touvron21a,
title = {Training data-efficient image transformers & distillation through attention},
author = {Touvron, Hugo and Cord, Matthieu and Douze, Matthijs and Massa, Francisco and Sablayrolles, Alexandre and Jegou, Herve},
booktitle = {International Conference on Machine Learning},
pages = {10347--10357},
year = {2021},
volume = {139},
month = {July}
}@misc{rw2019timm,
author = {Ross Wightman},
title = {PyTorch Image Models},
year = {2019},
publisher = {GitHub},
journal = {GitHub repository},
doi = {10.5281/zenodo.4414861},
howpublished = {\url{https://github.com/huggingface/pytorch-image-models}}
}