# coding=utf-8
# Copyright 2022 The IDEA Authors. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import Any, Dict
import torchvision
from detectron2.modeling.backbone import Backbone
try:
from torchvision.models.feature_extraction import (
create_feature_extractor,
)
has_feature_extractor = True
except ImportError:
has_feature_extractor = False
[docs]class TorchvisionBackbone(Backbone):
"""A wrapper for torchvision pretrained backbones
Please check `Feature extraction for model inspection
<https://pytorch.org/vision/stable/feature_extraction.html>`_
for more details.
Args:
model_name (str): Name of torchvision models. Default: resnet50.
pretrained (bool): Whether to load pretrained weights. Default: False.
weights (Optional[ResNet50_Weights]): The pretrained weights to use. Default: None.
return_nodes (Dict[str, str]): The keys are the node names and the values are the
user-specified keys for the graph module's returned dictionary.
"""
def __init__(
self,
model_name: str = "resnet50",
pretrained: bool = False,
return_nodes: Dict[str, str] = {
"layer1": "res2",
"layer2": "res3",
"layer3": "res4",
"layer4": "res5",
},
train_return_nodes: Dict[str, str] = None,
eval_return_nodes: Dict[str, str] = None,
tracer_kwargs: Dict[str, Any] = None,
suppress_diff_warnings: bool = False,
**kwargs,
):
super(TorchvisionBackbone, self).__init__()
# build torchvision models
self.model = getattr(torchvision.models, model_name)(pretrained=pretrained, **kwargs)
if has_feature_extractor is False:
raise RuntimeError(
"Failed to import create_feature_extractor from torchvision. \
Please install torchvision 1.10+."
)
# turn models into feature extractor
self.feature_extractor = create_feature_extractor(
model=self.model,
return_nodes=return_nodes,
train_return_nodes=train_return_nodes,
eval_return_nodes=eval_return_nodes,
tracer_kwargs=tracer_kwargs,
suppress_diff_warning=suppress_diff_warnings,
)
[docs] def forward(self, x):
"""Forward function of TorchvisionBackbone
Args:
x (torch.Tensor): the input tensor for feature extraction.
Returns:
dict[str->Tensor]: mapping from feature name (e.g., "res2") to tensor
"""
outs = self.feature_extractor(x)
return outs