File size: 3,060 Bytes
5f0437a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 |
import numpy as np
import torch.nn as nn
import torch
from torch.nn.modules import module
import torch.nn.functional as F
class MLP(nn.Module):
"""
Linear Embedding:
"""
def __init__(self, input_dim=2048, embed_dim=768):
super().__init__()
self.proj = nn.Linear(input_dim, embed_dim)
def forward(self, x):
x = x.flatten(2).transpose(1, 2)
x = self.proj(x)
return x
class DecoderHead(nn.Module):
def __init__(self,
in_channels=[64, 128, 320, 512],
num_classes=40,
dropout_ratio=0.1,
norm_layer=nn.BatchNorm2d,
embed_dim=768,
align_corners=False):
super(DecoderHead, self).__init__()
self.num_classes = num_classes
self.dropout_ratio = dropout_ratio
self.align_corners = align_corners
self.in_channels = in_channels
if dropout_ratio > 0:
self.dropout = nn.Dropout2d(dropout_ratio)
else:
self.dropout = None
c1_in_channels, c2_in_channels, c3_in_channels, c4_in_channels = self.in_channels
embedding_dim = embed_dim
self.linear_c4 = MLP(input_dim=c4_in_channels, embed_dim=embedding_dim)
self.linear_c3 = MLP(input_dim=c3_in_channels, embed_dim=embedding_dim)
self.linear_c2 = MLP(input_dim=c2_in_channels, embed_dim=embedding_dim)
self.linear_c1 = MLP(input_dim=c1_in_channels, embed_dim=embedding_dim)
self.linear_fuse = nn.Sequential(
nn.Conv2d(in_channels=embedding_dim*4, out_channels=embedding_dim, kernel_size=1),
norm_layer(embedding_dim),
nn.ReLU(inplace=True)
)
self.linear_pred = nn.Conv2d(embedding_dim, self.num_classes, kernel_size=1)
def forward(self, inputs, return_feats=False):
# len=4, 1/4,1/8,1/16,1/32
c1, c2, c3, c4 = inputs
############## MLP decoder on C1-C4 ###########
n, _, h, w = c4.shape
_c4 = self.linear_c4(c4).permute(0,2,1).reshape(n, -1, c4.shape[2], c4.shape[3])
_c4 = F.interpolate(_c4, size=c1.size()[2:],mode='bilinear',align_corners=self.align_corners)
_c3 = self.linear_c3(c3).permute(0,2,1).reshape(n, -1, c3.shape[2], c3.shape[3])
_c3 = F.interpolate(_c3, size=c1.size()[2:],mode='bilinear',align_corners=self.align_corners)
_c2 = self.linear_c2(c2).permute(0,2,1).reshape(n, -1, c2.shape[2], c2.shape[3])
_c2 = F.interpolate(_c2, size=c1.size()[2:],mode='bilinear',align_corners=self.align_corners)
_c1 = self.linear_c1(c1).permute(0,2,1).reshape(n, -1, c1.shape[2], c1.shape[3])
_c = torch.cat([_c4, _c3, _c2, _c1], dim=1)
x = self.linear_fuse(_c)
x = self.dropout(x)
x = self.linear_pred(x)
if return_feats:
return x, _c
else:
return x
|