| 22 | |
| 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 |
| 64 | c1, c2, c3, c4 = inputs |
| 65 | |
| 66 | ############## MLP decoder on C1-C4 ########### |
| 67 | n, _, h, w = c4.shape |
| 68 | |
| 69 | _c4 = self.linear_c4(c4).permute(0, 2, 1).reshape(n, -1, c4.shape[2], c4.shape[3]) |
| 70 | _c4 = F.interpolate(_c4, size=c1.size()[2:], mode="bilinear", align_corners=self.align_corners) |
| 71 | |
| 72 | _c3 = self.linear_c3(c3).permute(0, 2, 1).reshape(n, -1, c3.shape[2], c3.shape[3]) |
| 73 | _c3 = F.interpolate(_c3, size=c1.size()[2:], mode="bilinear", align_corners=self.align_corners) |
| 74 | |
| 75 | _c2 = self.linear_c2(c2).permute(0, 2, 1).reshape(n, -1, c2.shape[2], c2.shape[3]) |
| 76 | _c2 = F.interpolate(_c2, size=c1.size()[2:], mode="bilinear", align_corners=self.align_corners) |
| 77 | |
| 78 | _c1 = self.linear_c1(c1).permute(0, 2, 1).reshape(n, -1, c1.shape[2], c1.shape[3]) |
| 79 | |
| 80 | _c = self.linear_fuse(torch.cat([_c4, _c3, _c2, _c1], dim=1)) |
| 81 | x = self.dropout(_c) |