|
|
| 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
|
| }
|
|
|
