File size: 1,540 Bytes
b13b124 |
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 |
import torch
import torch.nn as nn
import torch.nn.functional as F
from functools import partial
import math
from .helpers import load_pretrained
from .layers import DropPath, to_2tuple, trunc_normal_
from ..builder import HEADS
from .decode_head import BaseDecodeHead
from ..backbones.vit import Block
from mmcv.cnn import build_norm_layer
@HEADS.register_module()
class VIT_MLA_AUXIHead(BaseDecodeHead):
""" Vision Transformer with support for patch or hybrid CNN input stage
"""
def __init__(self, img_size=768, **kwargs):
super(VIT_MLA_AUXIHead, self).__init__(**kwargs)
self.img_size = img_size
if self.in_channels==1024:
self.aux_0 = nn.Conv2d(self.in_channels, 256, kernel_size=1, bias=False)
self.aux_1 = nn.Conv2d(256, self.num_classes, kernel_size=1, bias=False)
elif self.in_channels==256:
self.aux = nn.Conv2d(self.in_channels, self.num_classes, kernel_size=1, bias=False)
def to_2D(self, x):
n, hw, c = x.shape
h=w = int(math.sqrt(hw))
x = x.transpose(1,2).reshape(n, c, h, w)
return x
def forward(self, x):
x = self._transform_inputs(x)
if x.dim()==3:
x = x[:,1:]
x = self.to_2D(x)
if self.in_channels==1024:
x = self.aux_0(x)
x = self.aux_1(x)
elif self.in_channels==256:
x = self.aux(x)
x = F.interpolate(x, size=self.img_size, mode='bilinear', align_corners=self.align_corners)
return x
|