Skip to content
GitLab
Menu
Projects
Groups
Snippets
Loading...
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Contribute to GitLab
Sign in / Register
Toggle navigation
Menu
Open sidebar
OpenDAS
vision
Commits
c5e972a6
Commit
c5e972a6
authored
Jan 16, 2020
by
Lara Haidar
Committed by
Francisco Massa
Jan 16, 2020
Browse files
UIpdate Doc with ONNX support (#1752)
* update doc * update doc
parent
add2ecc7
Changes
3
Hide whitespace changes
Inline
Side-by-side
Showing
3 changed files
with
15 additions
and
0 deletions
+15
-0
torchvision/models/detection/faster_rcnn.py
torchvision/models/detection/faster_rcnn.py
+5
-0
torchvision/models/detection/keypoint_rcnn.py
torchvision/models/detection/keypoint_rcnn.py
+5
-0
torchvision/models/detection/mask_rcnn.py
torchvision/models/detection/mask_rcnn.py
+5
-0
No files found.
torchvision/models/detection/faster_rcnn.py
View file @
c5e972a6
...
@@ -315,6 +315,8 @@ def fasterrcnn_resnet50_fpn(pretrained=False, progress=True,
...
@@ -315,6 +315,8 @@ def fasterrcnn_resnet50_fpn(pretrained=False, progress=True,
- labels (``Int64Tensor[N]``): the predicted labels for each image
- labels (``Int64Tensor[N]``): the predicted labels for each image
- scores (``Tensor[N]``): the scores or each prediction
- scores (``Tensor[N]``): the scores or each prediction
Faster R-CNN is exportable to ONNX for a fixed batch size with inputs images of fixed size.
Example::
Example::
>>> model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
>>> model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
...
@@ -333,6 +335,9 @@ def fasterrcnn_resnet50_fpn(pretrained=False, progress=True,
...
@@ -333,6 +335,9 @@ def fasterrcnn_resnet50_fpn(pretrained=False, progress=True,
>>> model.eval()
>>> model.eval()
>>> x = [torch.rand(3, 300, 400), torch.rand(3, 500, 400)]
>>> x = [torch.rand(3, 300, 400), torch.rand(3, 500, 400)]
>>> predictions = model(x)
>>> predictions = model(x)
>>>
>>> # optionally, if you want to export the model to ONNX:
>>> torch.onnx.export(model, x, "faster_rcnn.onnx", opset_version = 11)
Arguments:
Arguments:
pretrained (bool): If True, returns a model pre-trained on COCO train2017
pretrained (bool): If True, returns a model pre-trained on COCO train2017
...
...
torchvision/models/detection/keypoint_rcnn.py
View file @
c5e972a6
...
@@ -298,12 +298,17 @@ def keypointrcnn_resnet50_fpn(pretrained=False, progress=True,
...
@@ -298,12 +298,17 @@ def keypointrcnn_resnet50_fpn(pretrained=False, progress=True,
- scores (``Tensor[N]``): the scores or each prediction
- scores (``Tensor[N]``): the scores or each prediction
- keypoints (``FloatTensor[N, K, 3]``): the locations of the predicted keypoints, in ``[x, y, v]`` format.
- keypoints (``FloatTensor[N, K, 3]``): the locations of the predicted keypoints, in ``[x, y, v]`` format.
Keypoint R-CNN is exportable to ONNX for a fixed batch size with inputs images of fixed size.
Example::
Example::
>>> model = torchvision.models.detection.keypointrcnn_resnet50_fpn(pretrained=True)
>>> model = torchvision.models.detection.keypointrcnn_resnet50_fpn(pretrained=True)
>>> model.eval()
>>> model.eval()
>>> x = [torch.rand(3, 300, 400), torch.rand(3, 500, 400)]
>>> x = [torch.rand(3, 300, 400), torch.rand(3, 500, 400)]
>>> predictions = model(x)
>>> predictions = model(x)
>>>
>>> # optionally, if you want to export the model to ONNX:
>>> torch.onnx.export(model, x, "keypoint_rcnn.onnx", opset_version = 11)
Arguments:
Arguments:
pretrained (bool): If True, returns a model pre-trained on COCO train2017
pretrained (bool): If True, returns a model pre-trained on COCO train2017
...
...
torchvision/models/detection/mask_rcnn.py
View file @
c5e972a6
...
@@ -294,12 +294,17 @@ def maskrcnn_resnet50_fpn(pretrained=False, progress=True,
...
@@ -294,12 +294,17 @@ def maskrcnn_resnet50_fpn(pretrained=False, progress=True,
obtain the final segmentation masks, the soft masks can be thresholded, generally
obtain the final segmentation masks, the soft masks can be thresholded, generally
with a value of 0.5 (``mask >= 0.5``)
with a value of 0.5 (``mask >= 0.5``)
Mask R-CNN is exportable to ONNX for a fixed batch size with inputs images of fixed size.
Example::
Example::
>>> model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True)
>>> model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True)
>>> model.eval()
>>> model.eval()
>>> x = [torch.rand(3, 300, 400), torch.rand(3, 500, 400)]
>>> x = [torch.rand(3, 300, 400), torch.rand(3, 500, 400)]
>>> predictions = model(x)
>>> predictions = model(x)
>>>
>>> # optionally, if you want to export the model to ONNX:
>>> torch.onnx.export(model, x, "mask_rcnn.onnx", opset_version = 11)
Arguments:
Arguments:
pretrained (bool): If True, returns a model pre-trained on COCO train2017
pretrained (bool): If True, returns a model pre-trained on COCO train2017
...
...
Write
Preview
Markdown
is supported
0%
Try again
or
attach a new file
.
Attach a file
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment