一个 HRNet 图像分类模型。由论文作者在 ImageNet-1k 上训练。
from PIL import Image
import timm
from timm.models.efficientnet import _cfg
from openmind import is_torch_npu_available
import torch_npu
import torch
import argparse
from openmind_hub import snapshot_download
if is_torch_npu_available():
device = "npu:0"
else:
device = "cpu"
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument( "--model_name_or_path", type=str, default="hrnet_w18.ms_aug_in1k", )
args = parser.parse_args()
return args
args = parse_args()
if args.model_name_or_path:
model_path = args.model_name_or_path
else:
model_path = snapshot_download(
"CICC/hrnet_w18.ms_aug_in1k",
revision="main",
resume_download=True,
ignore_patterns=["*.h5", "*.ot", " *.msgpack"]
)
# load tokenizer
img = Image.open('./beignets-task-guide.png')
config = _cfg(url='', file='model.safetensors')
model = timm.create_model("hrnet_w18.ms_aug_in1k", pretrained=True, pretrained_cfg=config).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).npu()) # unsqueeze single image into batch of 1
print(output)
for o in output:
# print shape of each feature map in output
# e.g.:
# torch.Size([1, 64, 147, 147])
# torch.Size([1, 192, 71, 71])
# torch.Size([1, 288, 35, 35])
# torch.Size([1, 768, 17, 17])
# torch.Size([1, 2048, 8, 8])
print(o.shape)在 timm 模型结果 中探索此模型的数据集和运行时指标。
@article{WangSCJDZLMTWLX19,
title={Deep High-Resolution Representation Learning for Visual Recognition},
author={Jingdong Wang and Ke Sun and Tianheng Cheng and
Borui Jiang and Chaorui Deng and Yang Zhao and Dong Liu and Yadong Mu and
Mingkui Tan and Xinggang Wang and Wenyu Liu and Bin Xiao},
journal = {TPAMI}
year={2019}
}