# 2022.06.08-Changed for implementation of TokenFusion # Huawei Technologies Co., Ltd. import torch import torch.nn as nn import torch.nn.functional as F from . import mix_transformer from mmcv.cnn import ConvModule from .modules import num_parallel 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).contiguous() x = self.proj(x) return x class SegFormerHead(nn.Module): """ SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers """ def __init__(self, feature_strides=None, in_channels=128, embedding_dim=256, num_classes=20, **kwargs): super(SegFormerHead, self).__init__() self.in_channels = in_channels self.num_classes = num_classes assert len(feature_strides) == len(self.in_channels) assert min(feature_strides) == feature_strides[0] self.feature_strides = feature_strides c1_in_channels, c2_in_channels, c3_in_channels, c4_in_channels = self.in_channels #decoder_params = kwargs['decoder_params'] #embedding_dim = decoder_params['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.dropout = nn.Dropout2d(0.1) self.linear_fuse = ConvModule( in_channels=embedding_dim*4, out_channels=embedding_dim, kernel_size=1, norm_cfg=dict(type='BN', requires_grad=True) ) self.linear_pred = nn.Conv2d(embedding_dim, self.num_classes, kernel_size=1) def forward(self, x): c1, c2, c3, c4 = x ############## 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]).contiguous() _c4 = F.interpolate(_c4, size=c1.size()[2:],mode='bilinear',align_corners=False) _c3 = self.linear_c3(c3).permute(0,2,1).reshape(n, -1, c3.shape[2], c3.shape[3]).contiguous() _c3 = F.interpolate(_c3, size=c1.size()[2:],mode='bilinear',align_corners=False) _c2 = self.linear_c2(c2).permute(0,2,1).reshape(n, -1, c2.shape[2], c2.shape[3]).contiguous() _c2 = F.interpolate(_c2, size=c1.size()[2:],mode='bilinear',align_corners=False) _c1 = self.linear_c1(c1).permute(0,2,1).reshape(n, -1, c1.shape[2], c1.shape[3]).contiguous() _c = self.linear_fuse(torch.cat([_c4, _c3, _c2, _c1], dim=1)) x = self.dropout(_c) x = self.linear_pred(x) return x class UnifiedRepresentationNetwork(nn.Module): def __init__(self, backbone, config, num_classes=20, embedding_dim=256, pretrained=True): super().__init__() self.num_classes = num_classes self.embedding_dim = embedding_dim self.feature_strides = [4, 8, 16, 32] self.num_parallel = num_parallel #self.in_channels = [32, 64, 160, 256] #self.in_channels = [64, 128, 320, 512] self.encoder0 = getattr(mix_transformer, backbone)() self.encoder1 = getattr(mix_transformer, backbone)() self.in_channels = self.encoder0.embed_dims ## initilize encoder if pretrained: state_dict0 = torch.load(config.root_dir+'/data/pytorch-weight/' + backbone + '.pth') state_dict0.pop('head.weight') state_dict0.pop('head.bias') state_dict0 = expand_state_dict(self.encoder0.state_dict(), state_dict0, self.num_parallel) self.encoder0.load_state_dict(state_dict0, strict=True) state_dict1 = torch.load(config.root_dir+'/data/pytorch-weight/' + backbone + '.pth') state_dict1.pop('head.weight') state_dict1.pop('head.bias') state_dict1 = expand_state_dict(self.encoder1.state_dict(), state_dict1, self.num_parallel) self.encoder1.load_state_dict(state_dict1, strict=True) print("Load pretrained weights from " + config.root_dir+'/data/pytorch-weight/' + backbone + '.pth') else: print("Train from scratch") self.decoder = SegFormerHead(feature_strides=self.feature_strides, in_channels=self.in_channels, embedding_dim=self.embedding_dim, num_classes=self.num_classes) # self.alpha = nn.Parameter(torch.ones(self.num_parallel, requires_grad=True)) # self.register_parameter('alpha', self.alpha) def get_params(self): param_groups = [[], [], []] for name, param in list(self.encoder0.named_parameters()): if "norm" in name: param_groups[1].append(param) else: param_groups[0].append(param) for name, param in list(self.encoder1.named_parameters()): if "norm" in name: param_groups[1].append(param) else: param_groups[0].append(param) for param in list(self.decoder.parameters()): param_groups[2].append(param) return param_groups # def get_params(self): # param_groups = [[], []] # for param in list(self.encoder.parameters()): # param_groups[0].append(param) # for param in list(self.decoder.parameters()): # param_groups[1].append(param) # return param_groups def forward(self, data, get_sup_loss = False, gt = None, criterion = None): b, c, h, w = data[0].shape #rgb is the 0th element x0 = self.encoder0([data[0]])[0] x1 = self.encoder1([data[1]])[0] encoded = [] for enc0, enc1 in zip(x0, x1): # encoded.append(enc1) encoded.append((enc0 + enc1) / 2) pred = [self.decoder(encoded)] for i in range(len(pred)): pred[i] = F.interpolate(pred[i], size=(h, w), mode='bilinear', align_corners=True) if not self.training: return pred else: # training if get_sup_loss: sup_loss = self.get_sup_loss(pred, gt, criterion) # print(sup_loss, l1, l1_loss, sup_loss + l1_loss, "losses") return pred, sup_loss else: return pred def get_sup_loss(self, pred, gt, criterion): sup_loss = 0 for p in pred: p = p[:gt.shape[0]] #Getting loss for only those examples in batch where gt exists. Won't get sup loss for unlabeled data. # soft_output = nn.LogSoftmax()(p) sup_loss += criterion(p, gt) return sup_loss / len(pred) def expand_state_dict(model_dict, state_dict, num_parallel): model_dict_keys = model_dict.keys() state_dict_keys = state_dict.keys() for model_dict_key in model_dict_keys: model_dict_key_re = model_dict_key.replace('module.', '') if model_dict_key_re in state_dict_keys: model_dict[model_dict_key] = state_dict[model_dict_key_re] for i in range(num_parallel): ln = '.ln_%d' % i replace = True if ln in model_dict_key_re else False model_dict_key_re = model_dict_key_re.replace(ln, '') if replace and model_dict_key_re in state_dict_keys: model_dict[model_dict_key] = state_dict[model_dict_key_re] return model_dict