| 767 | |
| 768 | |
| 769 | class D2SwinTransformer(SwinTransformer, Backbone): |
| 770 | def __init__(self, cfg, pretrain_img_size, patch_size, in_chans, embed_dim, |
| 771 | depths, num_heads, window_size, mlp_ratio, qkv_bias, qk_scale, |
| 772 | drop_rate, attn_drop_rate, drop_path_rate, norm_layer, ape, |
| 773 | patch_norm, out_indices, use_checkpoint): |
| 774 | super().__init__( |
| 775 | pretrain_img_size, |
| 776 | patch_size, |
| 777 | in_chans, |
| 778 | embed_dim, |
| 779 | depths, |
| 780 | num_heads, |
| 781 | window_size, |
| 782 | mlp_ratio, |
| 783 | qkv_bias, |
| 784 | qk_scale, |
| 785 | drop_rate, |
| 786 | attn_drop_rate, |
| 787 | drop_path_rate, |
| 788 | norm_layer, |
| 789 | ape, |
| 790 | patch_norm, |
| 791 | out_indices, |
| 792 | use_checkpoint=use_checkpoint, |
| 793 | ) |
| 794 | |
| 795 | self._out_features = cfg['OUT_FEATURES'] |
| 796 | |
| 797 | self._out_feature_strides = { |
| 798 | "res2": 4, |
| 799 | "res3": 8, |
| 800 | "res4": 16, |
| 801 | "res5": 32, |
| 802 | } |
| 803 | self._out_feature_channels = { |
| 804 | "res2": self.num_features[0], |
| 805 | "res3": self.num_features[1], |
| 806 | "res4": self.num_features[2], |
| 807 | "res5": self.num_features[3], |
| 808 | } |
| 809 | |
| 810 | def forward(self, x): |
| 811 | """ |
| 812 | Args: |
| 813 | x: Tensor of shape (N,C,H,W). H, W must be a multiple of ``self.size_divisibility``. |
| 814 | Returns: |
| 815 | dict[str->Tensor]: names and the corresponding features |
| 816 | """ |
| 817 | assert ( |
| 818 | x.dim() == 4 |
| 819 | ), f"SwinTransformer takes an input of shape (N, C, H, W). Got {x.shape} instead!" |
| 820 | outputs = {} |
| 821 | y = super().forward(x) |
| 822 | for k in y.keys(): |
| 823 | if k in self._out_features: |
| 824 | outputs[k] = y[k] |
| 825 | return outputs |
| 826 | |