File size: 1,532 Bytes
b84549f |
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 |
from torch import nn
from einops import rearrange
import torch.nn.functional as F
from utils.dl.common.model import get_super_module
class DecoderLinear(nn.Module):
def __init__(self, n_cls, patch_size, d_encoder, im_size):
super(DecoderLinear, self).__init__()
self.d_encoder = d_encoder
self.patch_size = patch_size
self.n_cls = n_cls
self.im_size = im_size
self.head = nn.Linear(self.d_encoder, n_cls)
def debug(self):
print(self.head, id(self), 'debug()')
def forward(self, x):
# print('inside debug')
# self.debug()
x = x[:, 1:] # remove cls token
# print(x.size())
H, W = self.im_size
GS = H // self.patch_size
# print(H, W, GS, self.patch_size)
# print('head', self.head.weight.size(), x.size())
# print(self.head, 'debug()')
x = self.head(x)
# print(x.size())
# (b, HW//ps**2, ps_c)
x = rearrange(x, "b (h w) c -> b c h w", h=GS)
# print(x.size())
masks = x
masks = F.upsample(masks, size=(H, W), mode="bilinear")
# print(masks.size())
return masks
def modify_forward_head():
from types import MethodType
from timm.models.vision_transformer import VisionTransformer
def forward_head(self, x, pre_logits: bool = False):
return self.head(x)
VisionTransformer.forward_head = MethodType(forward_head, VisionTransformer)
|