| 23 | |
| 24 | class DecoderHead(nn.Module): |
| 25 | def __init__( |
| 26 | self, |
| 27 | in_channels=[64, 128, 320, 512], |
| 28 | num_classes=40, |
| 29 | dropout_ratio=0.1, |
| 30 | norm_layer=nn.BatchNorm2d, |
| 31 | embed_dim=768, |
| 32 | align_corners=False, |
| 33 | ): |
| 34 | super(DecoderHead, self).__init__() |
| 35 | self.num_classes = num_classes |
| 36 | self.dropout_ratio = dropout_ratio |
| 37 | self.align_corners = align_corners |
| 38 | |
| 39 | self.in_channels = in_channels |
| 40 | |
| 41 | if dropout_ratio > 0: |
| 42 | self.dropout = nn.Dropout2d(dropout_ratio) |
| 43 | else: |
| 44 | self.dropout = None |
| 45 | |
| 46 | c1_in_channels, c2_in_channels, c3_in_channels, c4_in_channels = self.in_channels |
| 47 | |
| 48 | embedding_dim = embed_dim |
| 49 | self.linear_c4 = MLP(input_dim=c4_in_channels, embed_dim=embedding_dim) |
| 50 | self.linear_c3 = MLP(input_dim=c3_in_channels, embed_dim=embedding_dim) |
| 51 | self.linear_c2 = MLP(input_dim=c2_in_channels, embed_dim=embedding_dim) |
| 52 | self.linear_c1 = MLP(input_dim=c1_in_channels, embed_dim=embedding_dim) |
| 53 | |
| 54 | self.linear_fuse = nn.Sequential( |
| 55 | nn.Conv2d(in_channels=embedding_dim * 4, out_channels=embedding_dim, kernel_size=1), |
| 56 | norm_layer(embedding_dim), |
| 57 | nn.ReLU(inplace=True), |
| 58 | ) |
| 59 | |
| 60 | self.linear_pred = nn.Conv2d(embedding_dim, self.num_classes, kernel_size=1) |
| 61 | |
| 62 | def forward(self, inputs): |
| 63 | # len=4, 1/4,1/8,1/16,1/32 |