|
|
|
from abc import ABCMeta, abstractmethod |
|
from typing import Dict |
|
import torch.nn as nn |
|
|
|
from detectron2.layers import ShapeSpec |
|
|
|
__all__ = ["Backbone"] |
|
|
|
|
|
class Backbone(nn.Module, metaclass=ABCMeta): |
|
""" |
|
Abstract base class for network backbones. |
|
""" |
|
|
|
def __init__(self): |
|
""" |
|
The `__init__` method of any subclass can specify its own set of arguments. |
|
""" |
|
super().__init__() |
|
|
|
@abstractmethod |
|
def forward(self): |
|
""" |
|
Subclasses must override this method, but adhere to the same return type. |
|
|
|
Returns: |
|
dict[str->Tensor]: mapping from feature name (e.g., "res2") to tensor |
|
""" |
|
pass |
|
|
|
@property |
|
def size_divisibility(self) -> int: |
|
""" |
|
Some backbones require the input height and width to be divisible by a |
|
specific integer. This is typically true for encoder / decoder type networks |
|
with lateral connection (e.g., FPN) for which feature maps need to match |
|
dimension in the "bottom up" and "top down" paths. Set to 0 if no specific |
|
input size divisibility is required. |
|
""" |
|
return 0 |
|
|
|
@property |
|
def padding_constraints(self) -> Dict[str, int]: |
|
""" |
|
This property is a generalization of size_divisibility. Some backbones and training |
|
recipes require specific padding constraints, such as enforcing divisibility by a specific |
|
integer (e.g., FPN) or padding to a square (e.g., ViTDet with large-scale jitter |
|
in :paper:vitdet). `padding_constraints` contains these optional items like: |
|
{ |
|
"size_divisibility": int, |
|
"square_size": int, |
|
# Future options are possible |
|
} |
|
`size_divisibility` will read from here if presented and `square_size` indicates the |
|
square padding size if `square_size` > 0. |
|
|
|
TODO: use type of Dict[str, int] to avoid torchscipt issues. The type of padding_constraints |
|
could be generalized as TypedDict (Python 3.8+) to support more types in the future. |
|
""" |
|
return {} |
|
|
|
def output_shape(self): |
|
""" |
|
Returns: |
|
dict[str->ShapeSpec] |
|
""" |
|
|
|
return { |
|
name: ShapeSpec( |
|
channels=self._out_feature_channels[name], stride=self._out_feature_strides[name] |
|
) |
|
for name in self._out_features |
|
} |
|
|
