|
1 |
| -from collections import OrderedDict |
2 |
| - |
3 | 1 | import torch
|
4 | 2 | from torch import nn
|
5 | 3 | import torch.nn.functional as F
|
6 | 4 |
|
7 |
| -from torchvision.ops import misc as misc_nn_ops |
8 | 5 | from torchvision.ops import MultiScaleRoIAlign
|
9 | 6 |
|
10 | 7 | from ._utils import overwrite_eps
|
|
19 | 16 |
|
20 | 17 |
|
21 | 18 | __all__ = [
|
22 |
| - "FasterRCNN", "fasterrcnn_resnet50_fpn", "fasterrcnn_mobilenet_v3_large", "fasterrcnn_mobilenet_v3_large_fpn" |
| 19 | + "FasterRCNN", "fasterrcnn_resnet50_fpn", "fasterrcnn_mobilenet_v3_large_fpn" |
23 | 20 | ]
|
24 | 21 |
|
25 | 22 |
|
@@ -291,8 +288,7 @@ def forward(self, x):
|
291 | 288 | model_urls = {
|
292 | 289 | 'fasterrcnn_resnet50_fpn_coco':
|
293 | 290 | 'https://download.pytorch.org/models/fasterrcnn_resnet50_fpn_coco-258fb6c6.pth',
|
294 |
| - 'fasterrcnn_mobilenet_v3_large_coco': None, |
295 |
| - 'fasterrcnn_mobilenet_v3_large_fpn_coco': None, |
| 291 | + 'fasterrcnn_mobilenet_v3_large_fpn_coco': None, # TODO: Add the final model url |
296 | 292 | }
|
297 | 293 |
|
298 | 294 |
|
@@ -371,48 +367,6 @@ def fasterrcnn_resnet50_fpn(pretrained=False, progress=True,
|
371 | 367 | return model
|
372 | 368 |
|
373 | 369 |
|
374 |
| -def fasterrcnn_mobilenet_v3_large(pretrained=False, progress=True, num_classes=91, pretrained_backbone=True, |
375 |
| - trainable_backbone_layers=None, min_size=320, max_size=640, **kwargs): |
376 |
| - """ |
377 |
| - Constructs a Faster R-CNN model with a MobileNetV3-Large backbone. It works similarly |
378 |
| - to Faster R-CNN with ResNet-50 FPN backbone. See `fasterrcnn_resnet50_fpn` for more details. |
379 |
| -
|
380 |
| - Example:: |
381 |
| -
|
382 |
| - >>> model = torchvision.models.detection.fasterrcnn_mobilenet_v3_large(pretrained=True) |
383 |
| - >>> model.eval() |
384 |
| - >>> x = [torch.rand(3, 300, 400), torch.rand(3, 500, 400)] |
385 |
| - >>> predictions = model(x) |
386 |
| -
|
387 |
| - Args: |
388 |
| - pretrained (bool): If True, returns a model pre-trained on COCO train2017 |
389 |
| - progress (bool): If True, displays a progress bar of the download to stderr |
390 |
| - num_classes (int): number of output classes of the model (including the background) |
391 |
| - pretrained_backbone (bool): If True, returns a model with backbone pre-trained on Imagenet |
392 |
| - trainable_backbone_layers (int): number of trainable (not frozen) resnet layers starting from final block. |
393 |
| - Valid values are between 0 and 6, with 6 meaning all backbone layers are trainable. |
394 |
| - min_size (int): minimum size of the image to be rescaled before feeding it to the backbone |
395 |
| - max_size (int): maximum size of the image to be rescaled before feeding it to the backbone |
396 |
| - """ |
397 |
| - trainable_backbone_layers = _validate_trainable_layers( |
398 |
| - pretrained or pretrained_backbone, trainable_backbone_layers, 6, 3) |
399 |
| - |
400 |
| - if pretrained: |
401 |
| - pretrained_backbone = False |
402 |
| - backbone = mobilenet_backbone("mobilenet_v3_large", pretrained_backbone, False, |
403 |
| - trainable_layers=trainable_backbone_layers) |
404 |
| - |
405 |
| - anchor_sizes = ((32, 64, 128, 256, 512, ), ) |
406 |
| - aspect_ratios = ((0.5, 1.0, 2.0), ) |
407 |
| - |
408 |
| - model = FasterRCNN(backbone, num_classes, rpn_anchor_generator=AnchorGenerator(anchor_sizes, aspect_ratios), |
409 |
| - min_size=min_size, max_size=max_size, **kwargs) |
410 |
| - if pretrained: |
411 |
| - state_dict = load_state_dict_from_url(model_urls['fasterrcnn_mobilenet_v3_large_coco'], progress=progress) |
412 |
| - model.load_state_dict(state_dict) |
413 |
| - return model |
414 |
| - |
415 |
| - |
416 | 370 | def fasterrcnn_mobilenet_v3_large_fpn(pretrained=False, progress=True, num_classes=91, pretrained_backbone=True,
|
417 | 371 | trainable_backbone_layers=None, min_size=320, max_size=640, rpn_score_thresh=0.05,
|
418 | 372 | **kwargs):
|
|
0 commit comments