This is the multi-scale encoder in detection models, also named as pixel decoder in segmentation models.
| 176 | |
| 177 | |
| 178 | class MaskDINOEncoder(nn.Module): |
| 179 | """ |
| 180 | This is the multi-scale encoder in detection models, also named as pixel decoder in segmentation models. |
| 181 | """ |
| 182 | @configurable |
| 183 | def __init__( |
| 184 | self, |
| 185 | input_shape: Dict[str, ShapeSpec], |
| 186 | *, |
| 187 | transformer_dropout: float, |
| 188 | transformer_nheads: int, |
| 189 | transformer_dim_feedforward: int, |
| 190 | transformer_enc_layers: int, |
| 191 | conv_dim: int, |
| 192 | mask_dim: int, |
| 193 | norm: Optional[Union[str, Callable]] = None, |
| 194 | # deformable transformer encoder args |
| 195 | transformer_in_features: List[str], |
| 196 | common_stride: int, |
| 197 | num_feature_levels: int, |
| 198 | total_num_feature_levels: int, |
| 199 | feature_order: str, |
| 200 | use_ckpt=False, |
| 201 | ): |
| 202 | """ |
| 203 | NOTE: this interface is experimental. |
| 204 | Args: |
| 205 | input_shape: shapes (channels and stride) of the input features |
| 206 | transformer_dropout: dropout probability in transformer |
| 207 | transformer_nheads: number of heads in transformer |
| 208 | transformer_dim_feedforward: dimension of feedforward network |
| 209 | transformer_enc_layers: number of transformer encoder layers |
| 210 | conv_dims: number of output channels for the intermediate conv layers. |
| 211 | mask_dim: number of output channels for the final conv layer. |
| 212 | norm (str or callable): normalization for all conv layers |
| 213 | num_feature_levels: feature scales used |
| 214 | total_num_feature_levels: total feautre scales used (include the downsampled features) |
| 215 | feature_order: 'low2high' or 'high2low', i.e., 'low2high' means low-resolution features are put in the first. |
| 216 | """ |
| 217 | super().__init__() |
| 218 | self.use_ckpt = use_ckpt |
| 219 | transformer_input_shape = { |
| 220 | k: v for k, v in input_shape.items() if k in transformer_in_features |
| 221 | } |
| 222 | # this is the input shape of pixel decoder |
| 223 | input_shape = sorted(input_shape.items(), key=lambda x: x[1].stride) |
| 224 | self.in_features = [k for k, v in input_shape] # starting from "res2" to "res5" |
| 225 | self.feature_strides = [v.stride for k, v in input_shape] |
| 226 | self.feature_channels = [v.channels for k, v in input_shape] |
| 227 | self.feature_order = feature_order |
| 228 | |
| 229 | if feature_order == "low2high": |
| 230 | transformer_input_shape = sorted(transformer_input_shape.items(), key=lambda x: -x[1].stride) |
| 231 | else: |
| 232 | transformer_input_shape = sorted(transformer_input_shape.items(), key=lambda x: x[1].stride) |
| 233 | self.transformer_in_features = [k for k, v in transformer_input_shape] # starting from "res2" to "res5" |
| 234 | transformer_in_channels = [v.channels for k, v in transformer_input_shape] |
| 235 | self.transformer_feature_strides = [v.stride for k, v in transformer_input_shape] # to decide extra FPN layers |
no outgoing calls
no test coverage detected