| 16 | |
| 17 | |
| 18 | class TransformerDecoder(nn.Module): |
| 19 | |
| 20 | def __init__(self, decoder_layer, num_layers, norm=None, |
| 21 | return_intermediate=False, |
| 22 | d_model=256, query_dim=4, |
| 23 | modulate_hw_attn=True, |
| 24 | num_feature_levels=1, |
| 25 | deformable_decoder=True, |
| 26 | decoder_query_perturber=None, |
| 27 | dec_layer_number=None, # number of queries each layer in decoder |
| 28 | rm_dec_query_scale=True, |
| 29 | dec_layer_share=False, |
| 30 | dec_layer_dropout_prob=None, |
| 31 | task_switch=None, |
| 32 | ): |
| 33 | super().__init__() |
| 34 | if num_layers > 0: |
| 35 | self.layers = _get_clones(decoder_layer, num_layers, layer_share=dec_layer_share) |
| 36 | else: |
| 37 | self.layers = [] |
| 38 | self.num_layers = num_layers |
| 39 | self.norm = norm |
| 40 | self.return_intermediate = return_intermediate |
| 41 | assert return_intermediate, "support return_intermediate only" |
| 42 | self.query_dim = query_dim |
| 43 | assert query_dim in [2, 4], "query_dim should be 2/4 but {}".format(query_dim) |
| 44 | self.num_feature_levels = num_feature_levels |
| 45 | |
| 46 | self.ref_point_head = MLP(query_dim // 2 * d_model, d_model, d_model, 2) |
| 47 | if not deformable_decoder: |
| 48 | self.query_pos_sine_scale = MLP(d_model, d_model, d_model, 2) |
| 49 | else: |
| 50 | self.query_pos_sine_scale = None |
| 51 | |
| 52 | if rm_dec_query_scale: |
| 53 | self.query_scale = None |
| 54 | else: |
| 55 | raise NotImplementedError |
| 56 | self.query_scale = MLP(d_model, d_model, d_model, 2) |
| 57 | self.bbox_embed = None |
| 58 | self.class_embed = None |
| 59 | |
| 60 | self.d_model = d_model |
| 61 | self.modulate_hw_attn = modulate_hw_attn |
| 62 | self.deformable_decoder = deformable_decoder |
| 63 | |
| 64 | if not deformable_decoder and modulate_hw_attn: |
| 65 | self.ref_anchor_head = MLP(d_model, d_model, 2, 2) |
| 66 | else: |
| 67 | self.ref_anchor_head = None |
| 68 | |
| 69 | self.decoder_query_perturber = decoder_query_perturber |
| 70 | self.box_pred_damping = None |
| 71 | |
| 72 | self.dec_layer_number = dec_layer_number |
| 73 | if dec_layer_number is not None: |
| 74 | assert isinstance(dec_layer_number, list) |
| 75 | assert len(dec_layer_number) == num_layers |