mirrors
/
yolov7
mirror of https://github.com/WongKinYiu/yolov7.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

fix a model inference problem on cpu (#502) 

* fix model inference problem on cpu

* Update keypoint.ipynb
pull/530/head
spacewalk01 2022-08-16 22:23:16 +09:00 committed by GitHub
parent 36ce6b2087
commit 064c71e7c2
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 8 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
tools/keypoint.ipynb
View File

@ -25,10 +25,12 @@
"outputs": [],
"source": [
"device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
"weigths = torch.load('yolov7-w6-pose.pt')\n",
"weigths = torch.load('yolov7-w6-pose.pt', map_location=device)\n",
"model = weigths['model']\n",
"model = model.half().to(device)\n",
"_ = model.eval()"
"_ = model.float().eval()\n",
"\n",
"if torch.cuda.is_available():\n",
" model.half().to(device)"
]
},
{
@ -43,9 +45,9 @@
"image_ = image.copy()\n",
"image = transforms.ToTensor()(image)\n",
"image = torch.tensor(np.array([image.numpy()]))\n",
"image = image.to(device)\n",
"image = image.half()\n",
"\n",
"if torch.cuda.is_available():\n",
" image = image.half().to(device) \n",
"output, _ = model(image)"
]
},
@ -118,7 +120,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.8.10"
"version": "3.9.12"
}
},
"nbformat": 4,