|
print("Importing external...") |
|
import torch |
|
from torch import nn |
|
import torch.nn.functional as F |
|
|
|
from timm.models.efficientvit_mit import ( |
|
ConvNormAct, |
|
FusedMBConv, |
|
MBConv, |
|
ResidualBlock, |
|
efficientvit_l1, |
|
) |
|
from timm.layers import GELUTanh |
|
|
|
|
|
def val2list(x: list or tuple or any, repeat_time=1): |
|
if isinstance(x, (list, tuple)): |
|
return list(x) |
|
return [x for _ in range(repeat_time)] |
|
|
|
|
|
def resize( |
|
x: torch.Tensor, |
|
size: any or None = None, |
|
scale_factor: list[float] or None = None, |
|
mode: str = "bicubic", |
|
align_corners: bool or None = False, |
|
) -> torch.Tensor: |
|
if mode in {"bilinear", "bicubic"}: |
|
return F.interpolate( |
|
x, |
|
size=size, |
|
scale_factor=scale_factor, |
|
mode=mode, |
|
align_corners=align_corners, |
|
) |
|
elif mode in {"nearest", "area"}: |
|
return F.interpolate(x, size=size, scale_factor=scale_factor, mode=mode) |
|
else: |
|
raise NotImplementedError(f"resize(mode={mode}) not implemented.") |
|
|
|
|
|
class UpSampleLayer(nn.Module): |
|
def __init__( |
|
self, |
|
mode="bicubic", |
|
size: int or tuple[int, int] or list[int] or None = None, |
|
factor=2, |
|
align_corners=False, |
|
): |
|
super(UpSampleLayer, self).__init__() |
|
self.mode = mode |
|
self.size = val2list(size, 2) if size is not None else None |
|
self.factor = None if self.size is not None else factor |
|
self.align_corners = align_corners |
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor: |
|
if ( |
|
self.size is not None and tuple(x.shape[-2:]) == self.size |
|
) or self.factor == 1: |
|
return x |
|
return resize(x, self.size, self.factor, self.mode, self.align_corners) |
|
|
|
|
|
class DAGBlock(nn.Module): |
|
def __init__( |
|
self, |
|
inputs: dict[str, nn.Module], |
|
merge: str, |
|
post_input: nn.Module or None, |
|
middle: nn.Module, |
|
outputs: dict[str, nn.Module], |
|
): |
|
super(DAGBlock, self).__init__() |
|
|
|
self.input_keys = list(inputs.keys()) |
|
self.input_ops = nn.ModuleList(list(inputs.values())) |
|
self.merge = merge |
|
self.post_input = post_input |
|
|
|
self.middle = middle |
|
|
|
self.output_keys = list(outputs.keys()) |
|
self.output_ops = nn.ModuleList(list(outputs.values())) |
|
|
|
def forward(self, feature_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: |
|
feat = [ |
|
op(feature_dict[key]) for key, op in zip(self.input_keys, self.input_ops) |
|
] |
|
if self.merge == "add": |
|
feat = list_sum(feat) |
|
elif self.merge == "cat": |
|
feat = torch.concat(feat, dim=1) |
|
else: |
|
raise NotImplementedError |
|
if self.post_input is not None: |
|
feat = self.post_input(feat) |
|
feat = self.middle(feat) |
|
for key, op in zip(self.output_keys, self.output_ops): |
|
feature_dict[key] = op(feat) |
|
return feature_dict |
|
|
|
|
|
def list_sum(x: list) -> any: |
|
return x[0] if len(x) == 1 else x[0] + list_sum(x[1:]) |
|
|
|
|
|
class SegHead(nn.Module): |
|
def __init__( |
|
self, |
|
fid_list: list[str], |
|
in_channel_list: list[int], |
|
stride_list: list[int], |
|
head_stride: int, |
|
head_width: int, |
|
head_depth: int, |
|
expand_ratio: float, |
|
middle_op: str, |
|
final_expand: float or None, |
|
n_classes: int, |
|
dropout=0, |
|
norm="bn2d", |
|
act_func="hswish", |
|
): |
|
super(SegHead, self).__init__() |
|
|
|
if act_func == "gelu": |
|
act_func = GELUTanh |
|
else: |
|
raise ValueError(f"act_func {act_func} not supported") |
|
if norm == "bn2d": |
|
norm_layer = nn.BatchNorm2d |
|
else: |
|
raise ValueError(f"norm {norm} not supported") |
|
|
|
inputs = {} |
|
for fid, in_channel, stride in zip(fid_list, in_channel_list, stride_list): |
|
factor = stride // head_stride |
|
if factor == 1: |
|
inputs[fid] = ConvNormAct( |
|
in_channel, head_width, 1, norm_layer=norm_layer, act_layer=act_func |
|
) |
|
else: |
|
inputs[fid] = nn.Sequential( |
|
ConvNormAct( |
|
in_channel, |
|
head_width, |
|
1, |
|
norm_layer=norm_layer, |
|
act_layer=act_func, |
|
), |
|
UpSampleLayer(factor=factor), |
|
) |
|
self.in_keys = inputs.keys() |
|
self.in_ops = nn.ModuleList(inputs.values()) |
|
|
|
middle = [] |
|
for _ in range(head_depth): |
|
if middle_op == "mbconv": |
|
block = MBConv( |
|
head_width, |
|
head_width, |
|
expand_ratio=expand_ratio, |
|
norm_layer=norm_layer, |
|
act_layer=(act_func, act_func, None), |
|
) |
|
elif middle_op == "fmbconv": |
|
block = FusedMBConv( |
|
head_width, |
|
head_width, |
|
expand_ratio=expand_ratio, |
|
norm_layer=norm_layer, |
|
act_layer=(act_func, None), |
|
) |
|
else: |
|
raise NotImplementedError |
|
middle.append(ResidualBlock(block, nn.Identity())) |
|
self.middle = nn.Sequential(*middle) |
|
|
|
self.out_layer = nn.Sequential( |
|
*[ |
|
None |
|
if final_expand is None |
|
else ConvNormAct( |
|
head_width, |
|
head_width * final_expand, |
|
1, |
|
norm_layer=norm_layer, |
|
act_layer=act_func, |
|
), |
|
ConvNormAct( |
|
head_width * (final_expand or 1), |
|
n_classes, |
|
1, |
|
bias=True, |
|
dropout=dropout, |
|
norm_layer=None, |
|
act_layer=None, |
|
), |
|
] |
|
) |
|
|
|
def forward(self, feature_map_list): |
|
t_feat_maps = [ |
|
self.in_ops[ind](feature_map_list[ind]) |
|
for ind in range(len(feature_map_list)) |
|
] |
|
t_feat_map = list_sum(t_feat_maps) |
|
t_feat_map = self.middle(t_feat_map) |
|
out = self.out_layer(t_feat_map) |
|
return out |
|
|
|
|
|
class EfficientViT_l1_r224(nn.Module): |
|
def __init__( |
|
self, |
|
out_channels, |
|
out_ds_factor=1, |
|
decoder_size="small", |
|
pretrained=False, |
|
use_norm_params=False, |
|
): |
|
if decoder_size == "small": |
|
head_width = 32 |
|
head_depth = 1 |
|
middle_op = "mbconv" |
|
elif decoder_size == "medium": |
|
head_width = 64 |
|
head_depth = 3 |
|
middle_op = "mbconv" |
|
elif decoder_size == "large": |
|
head_width = 256 |
|
head_depth = 3 |
|
middle_op = "fmbconv" |
|
|
|
super(EfficientViT_l1_r224, self).__init__() |
|
self.bbone = efficientvit_l1( |
|
num_classes=0, features_only=True, pretrained=pretrained |
|
) |
|
self.head = SegHead( |
|
fid_list=["stage4", "stage3", "stage2"], |
|
in_channel_list=[512, 256, 128], |
|
stride_list=[32, 16, 8], |
|
head_stride=out_ds_factor, |
|
head_width=head_width, |
|
head_depth=head_depth, |
|
expand_ratio=4, |
|
middle_op=middle_op, |
|
final_expand=8, |
|
n_classes=out_channels, |
|
act_func="gelu", |
|
) |
|
|
|
if not use_norm_params: |
|
for module in self.modules(): |
|
if ( |
|
isinstance(module, nn.LayerNorm) |
|
or isinstance(module, nn.BatchNorm2d) |
|
or isinstance(module, nn.BatchNorm1d) |
|
): |
|
module.weight.requires_grad_(False) |
|
module.bias.requires_grad_(False) |
|
|
|
def forward(self, x): |
|
feat = self.bbone(x) |
|
out = self.head([feat[3], feat[2], feat[1]]) |
|
return out |
|
|
