Deeplabv3
Torchvision框架中在語(yǔ)義分割上支持的是Deeplabv3語(yǔ)義分割模型,而且支持不同的backbone替換,這些backbone替換包括MobileNetv3、ResNet50、ResNet101。其中MobileNetv3版本訓(xùn)練數(shù)據(jù)集是COCO子集,類別跟Pascal VOC的20個(gè)類別保持一致。這里以它為例,演示一下從模型導(dǎo)出ONNX到推理的全過(guò)程。ONNX格式導(dǎo)出
首先需要把pytorch的模型導(dǎo)出為onnx格式版本,用下面的腳本就好啦:
model=tv.models.segmentation.deeplabv3_mobilenet_v3_large(pretrained=True)
dummy_input=torch.randn(1,3,320,320)
model.eval()
model(dummy_input)
im=torch.zeros(1,3,320,320).to("cpu")
torch.onnx.export(model,im,
"deeplabv3_mobilenet.onnx",
verbose=False,
opset_version=11,
training=torch.onnx.TrainingMode.EVAL,
do_constant_folding=True,
input_names=['input'],
output_names=['out','aux'],
dynamic_axes={'input':{0:'batch',2:'height',3:'width'}}
)
模型的輸入與輸出結(jié)構(gòu)如下:
推理測(cè)試
模型推理對(duì)圖像有個(gè)預(yù)處理,要求如下:
transform = torchvision.transforms.Compose([
torchvision.transforms.ToTensor(),
torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
意思是轉(zhuǎn)換為0~1之間的浮點(diǎn)數(shù),然后減去均值除以方差。
剩下部分的代碼就比較簡(jiǎn)單,初始化onnx推理實(shí)例,然后完成推理,對(duì)結(jié)果完成解析,輸出推理結(jié)果,完整的代碼如下:
transform=torchvision.transforms.Compose([
torchvision.transforms.ToTensor(),
torchvision.transforms.Normalize(mean=[0.485,0.456,0.406],
std=[0.229,0.224,0.225])
])
sess_options=ort.SessionOptions()
#Belowisforoptimizingperformance
sess_options.intra_op_num_threads=24
#sess_options.execution_mode=ort.ExecutionMode.ORT_PARALLEL
sess_options.graph_optimization_level=ort.GraphOptimizationLevel.ORT_ENABLE_ALL
ort_session=ort.InferenceSession("deeplabv3_mobilenet.onnx",providers=['CUDAExecutionProvider'],sess_options=sess_options)
#src=cv.imread("D:/images/messi_player.jpg")
src=cv.imread("D:/images/master.jpg")
image=cv.cvtColor(src,cv.COLOR_BGR2RGB)
blob=transform(image)
c,h,w=blob.shape
input_x=blob.view(1,c,h,w)
defto_numpy(tensor):
returntensor.detach().cpu().numpy()iftensor.requires_gradelsetensor.cpu().numpy()
#computeONNXRuntimeoutputprediction
ort_inputs={ort_session.get_inputs()[0].name:to_numpy(input_x)}
ort_outs=ort_session.run(None,ort_inputs)
t1=ort_outs[0]
t2=ort_outs[1]
labels=np.argmax(np.squeeze(t1,0),axis=0)
print(labels.dtype,labels.shape)
red_map=np.zeros_like(labels).astype(np.uint8)
green_map=np.zeros_like(labels).astype(np.uint8)
blue_map=np.zeros_like(labels).astype(np.uint8)
forlabel_numinrange(0,len(label_color_map)):
index=labels==label_num
red_map[index]=np.array(label_color_map)[label_num,0]
green_map[index]=np.array(label_color_map)[label_num,1]
blue_map[index]=np.array(label_color_map)[label_num,2]
segmentation_map=np.stack([blue_map,green_map,red_map],axis=2)
cv.addWeighted(src,0.8,segmentation_map,0.2,0,src)
cv.imshow("deeplabv3",src)
cv.waitKey(0)
cv.destroyAllWindows()
運(yùn)行結(jié)果如下:
審核編輯 :李倩
-
數(shù)據(jù)集
+關(guān)注
關(guān)注
4文章
1224瀏覽量
25457 -
pytorch
+關(guān)注
關(guān)注
2文章
809瀏覽量
13967
原文標(biāo)題:輕松學(xué)Pytorch之Deeplabv3推理
文章出處:【微信號(hào):CVSCHOOL,微信公眾號(hào):OpenCV學(xué)堂】歡迎添加關(guān)注!文章轉(zhuǎn)載請(qǐng)注明出處。
發(fā)布評(píng)論請(qǐng)先 登錄
AI功能(SC171開發(fā)套件V3)
有獎(jiǎng)直播 | @4/8 輕松部署,強(qiáng)大擴(kuò)展邊緣運(yùn)算 AI 新世代

壁仞科技支持DeepSeek-V3滿血版訓(xùn)練推理
OpenAI即將推出o3 mini推理AI模型
OpenAI即將推出o3-mini,助力高級(jí)推理普及
利用Arm Kleidi技術(shù)實(shí)現(xiàn)PyTorch優(yōu)化

vLLM項(xiàng)目加入PyTorch生態(tài)系統(tǒng),引領(lǐng)LLM推理新紀(jì)元

Arm KleidiAI助力提升PyTorch上LLM推理性能

PyTorch 2.5.1: Bugs修復(fù)版發(fā)布

評(píng)論