You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
415 lines
9.6 KiB
Python
415 lines
9.6 KiB
Python
import torch
|
|
from torch import nn
|
|
|
|
from .vit import (
|
|
_make_pretrained_vitb16_384,
|
|
_make_pretrained_vitb_rn50_384,
|
|
_make_pretrained_vitl16_384,
|
|
)
|
|
|
|
|
|
def _make_encoder(
|
|
backbone,
|
|
features,
|
|
use_pretrained,
|
|
groups=1,
|
|
expand=False,
|
|
exportable=True,
|
|
hooks=None,
|
|
use_vit_only=False,
|
|
use_readout="ignore",
|
|
):
|
|
if backbone == "vitl16_384":
|
|
pretrained = _make_pretrained_vitl16_384(
|
|
use_pretrained, hooks=hooks, use_readout=use_readout
|
|
)
|
|
scratch = _make_scratch(
|
|
[256, 512, 1024, 1024], features, groups=groups, expand=expand
|
|
) # ViT-L/16 - 85.0% Top1 (backbone)
|
|
elif backbone == "vitb_rn50_384":
|
|
pretrained = _make_pretrained_vitb_rn50_384(
|
|
use_pretrained,
|
|
hooks=hooks,
|
|
use_vit_only=use_vit_only,
|
|
use_readout=use_readout,
|
|
)
|
|
scratch = _make_scratch(
|
|
[256, 512, 768, 768], features, groups=groups, expand=expand
|
|
) # ViT-H/16 - 85.0% Top1 (backbone)
|
|
elif backbone == "vitb16_384":
|
|
pretrained = _make_pretrained_vitb16_384(
|
|
use_pretrained, hooks=hooks, use_readout=use_readout
|
|
)
|
|
scratch = _make_scratch(
|
|
[96, 192, 384, 768], features, groups=groups, expand=expand
|
|
) # ViT-B/16 - 84.6% Top1 (backbone)
|
|
elif backbone == "resnext101_wsl":
|
|
pretrained = _make_pretrained_resnext101_wsl(use_pretrained)
|
|
scratch = _make_scratch(
|
|
[256, 512, 1024, 2048], features, groups=groups, expand=expand
|
|
) # efficientnet_lite3
|
|
elif backbone == "efficientnet_lite3":
|
|
pretrained = _make_pretrained_efficientnet_lite3(
|
|
use_pretrained, exportable=exportable
|
|
)
|
|
scratch = _make_scratch(
|
|
[32, 48, 136, 384], features, groups=groups, expand=expand
|
|
) # efficientnet_lite3
|
|
else:
|
|
msg = f"Backbone '{backbone}' not implemented"
|
|
raise NotImplementedError(msg)
|
|
|
|
return pretrained, scratch
|
|
|
|
|
|
def _make_scratch(in_shape, out_shape, groups=1, expand=False):
|
|
scratch = nn.Module()
|
|
|
|
out_shape1 = out_shape
|
|
out_shape2 = out_shape
|
|
out_shape3 = out_shape
|
|
out_shape4 = out_shape
|
|
if expand is True:
|
|
out_shape1 = out_shape
|
|
out_shape2 = out_shape * 2
|
|
out_shape3 = out_shape * 4
|
|
out_shape4 = out_shape * 8
|
|
|
|
scratch.layer1_rn = nn.Conv2d(
|
|
in_shape[0],
|
|
out_shape1,
|
|
kernel_size=3,
|
|
stride=1,
|
|
padding=1,
|
|
bias=False,
|
|
groups=groups,
|
|
)
|
|
scratch.layer2_rn = nn.Conv2d(
|
|
in_shape[1],
|
|
out_shape2,
|
|
kernel_size=3,
|
|
stride=1,
|
|
padding=1,
|
|
bias=False,
|
|
groups=groups,
|
|
)
|
|
scratch.layer3_rn = nn.Conv2d(
|
|
in_shape[2],
|
|
out_shape3,
|
|
kernel_size=3,
|
|
stride=1,
|
|
padding=1,
|
|
bias=False,
|
|
groups=groups,
|
|
)
|
|
scratch.layer4_rn = nn.Conv2d(
|
|
in_shape[3],
|
|
out_shape4,
|
|
kernel_size=3,
|
|
stride=1,
|
|
padding=1,
|
|
bias=False,
|
|
groups=groups,
|
|
)
|
|
|
|
return scratch
|
|
|
|
|
|
def _make_pretrained_efficientnet_lite3(use_pretrained, exportable=False):
|
|
efficientnet = torch.hub.load(
|
|
"rwightman/gen-efficientnet-pytorch",
|
|
"tf_efficientnet_lite3",
|
|
pretrained=use_pretrained,
|
|
exportable=exportable,
|
|
)
|
|
return _make_efficientnet_backbone(efficientnet)
|
|
|
|
|
|
def _make_efficientnet_backbone(effnet):
|
|
pretrained = nn.Module()
|
|
|
|
pretrained.layer1 = nn.Sequential(
|
|
effnet.conv_stem, effnet.bn1, effnet.act1, *effnet.blocks[0:2]
|
|
)
|
|
pretrained.layer2 = nn.Sequential(*effnet.blocks[2:3])
|
|
pretrained.layer3 = nn.Sequential(*effnet.blocks[3:5])
|
|
pretrained.layer4 = nn.Sequential(*effnet.blocks[5:9])
|
|
|
|
return pretrained
|
|
|
|
|
|
def _make_resnet_backbone(resnet):
|
|
pretrained = nn.Module()
|
|
pretrained.layer1 = nn.Sequential(
|
|
resnet.conv1, resnet.bn1, resnet.relu, resnet.maxpool, resnet.layer1
|
|
)
|
|
|
|
pretrained.layer2 = resnet.layer2
|
|
pretrained.layer3 = resnet.layer3
|
|
pretrained.layer4 = resnet.layer4
|
|
|
|
return pretrained
|
|
|
|
|
|
def _make_pretrained_resnext101_wsl(use_pretrained):
|
|
resnet = torch.hub.load("facebookresearch/WSL-Images", "resnext101_32x8d_wsl")
|
|
return _make_resnet_backbone(resnet)
|
|
|
|
|
|
class Interpolate(nn.Module):
|
|
"""Interpolation module."""
|
|
|
|
def __init__(self, scale_factor, mode, align_corners=False):
|
|
"""
|
|
Init.
|
|
|
|
Args:
|
|
scale_factor (float): scaling
|
|
mode (str): interpolation mode
|
|
"""
|
|
super().__init__()
|
|
|
|
self.interp = nn.functional.interpolate
|
|
self.scale_factor = scale_factor
|
|
self.mode = mode
|
|
self.align_corners = align_corners
|
|
|
|
def forward(self, x):
|
|
"""
|
|
Forward pass.
|
|
|
|
Args:
|
|
x (tensor): input
|
|
|
|
Returns:
|
|
tensor: interpolated data
|
|
"""
|
|
|
|
x = self.interp(
|
|
x,
|
|
scale_factor=self.scale_factor,
|
|
mode=self.mode,
|
|
align_corners=self.align_corners,
|
|
)
|
|
|
|
return x
|
|
|
|
|
|
class ResidualConvUnit(nn.Module):
|
|
"""Residual convolution module."""
|
|
|
|
def __init__(self, features):
|
|
"""
|
|
Init.
|
|
|
|
Args:
|
|
features (int): number of features
|
|
"""
|
|
super().__init__()
|
|
|
|
self.conv1 = nn.Conv2d(
|
|
features, features, kernel_size=3, stride=1, padding=1, bias=True
|
|
)
|
|
|
|
self.conv2 = nn.Conv2d(
|
|
features, features, kernel_size=3, stride=1, padding=1, bias=True
|
|
)
|
|
|
|
self.relu = nn.ReLU(inplace=True)
|
|
|
|
def forward(self, x):
|
|
"""
|
|
Forward pass.
|
|
|
|
Args:
|
|
x (tensor): input
|
|
|
|
Returns:
|
|
tensor: output
|
|
"""
|
|
out = self.relu(x)
|
|
out = self.conv1(out)
|
|
out = self.relu(out)
|
|
out = self.conv2(out)
|
|
|
|
return out + x
|
|
|
|
|
|
class FeatureFusionBlock(nn.Module):
|
|
"""Feature fusion block."""
|
|
|
|
def __init__(self, features):
|
|
"""
|
|
Init.
|
|
|
|
Args:
|
|
features (int): number of features
|
|
"""
|
|
super().__init__()
|
|
|
|
self.resConfUnit1 = ResidualConvUnit(features)
|
|
self.resConfUnit2 = ResidualConvUnit(features)
|
|
|
|
def forward(self, *xs):
|
|
"""
|
|
Forward pass.
|
|
|
|
Returns:
|
|
tensor: output
|
|
"""
|
|
output = xs[0]
|
|
|
|
if len(xs) == 2:
|
|
output += self.resConfUnit1(xs[1])
|
|
|
|
output = self.resConfUnit2(output)
|
|
|
|
output = nn.functional.interpolate(
|
|
output, scale_factor=2, mode="bilinear", align_corners=True
|
|
)
|
|
|
|
return output
|
|
|
|
|
|
class ResidualConvUnit_custom(nn.Module):
|
|
"""Residual convolution module."""
|
|
|
|
def __init__(self, features, activation, bn):
|
|
"""
|
|
Init.
|
|
|
|
Args:
|
|
features (int): number of features
|
|
"""
|
|
super().__init__()
|
|
|
|
self.bn = bn
|
|
|
|
self.groups = 1
|
|
|
|
self.conv1 = nn.Conv2d(
|
|
features,
|
|
features,
|
|
kernel_size=3,
|
|
stride=1,
|
|
padding=1,
|
|
bias=True,
|
|
groups=self.groups,
|
|
)
|
|
|
|
self.conv2 = nn.Conv2d(
|
|
features,
|
|
features,
|
|
kernel_size=3,
|
|
stride=1,
|
|
padding=1,
|
|
bias=True,
|
|
groups=self.groups,
|
|
)
|
|
|
|
if self.bn is True:
|
|
self.bn1 = nn.BatchNorm2d(features)
|
|
self.bn2 = nn.BatchNorm2d(features)
|
|
|
|
self.activation = activation
|
|
|
|
self.skip_add = nn.quantized.FloatFunctional()
|
|
|
|
def forward(self, x):
|
|
"""
|
|
Forward pass.
|
|
|
|
Args:
|
|
x (tensor): input
|
|
|
|
Returns:
|
|
tensor: output
|
|
"""
|
|
|
|
out = self.activation(x)
|
|
out = self.conv1(out)
|
|
if self.bn is True:
|
|
out = self.bn1(out)
|
|
|
|
out = self.activation(out)
|
|
out = self.conv2(out)
|
|
if self.bn is True:
|
|
out = self.bn2(out)
|
|
|
|
if self.groups > 1:
|
|
out = self.conv_merge(out)
|
|
|
|
return self.skip_add.add(out, x)
|
|
|
|
# return out + x
|
|
|
|
|
|
class FeatureFusionBlock_custom(nn.Module):
|
|
"""Feature fusion block."""
|
|
|
|
def __init__(
|
|
self,
|
|
features,
|
|
activation,
|
|
deconv=False,
|
|
bn=False,
|
|
expand=False,
|
|
align_corners=True,
|
|
):
|
|
"""
|
|
Init.
|
|
|
|
Args:
|
|
features (int): number of features
|
|
"""
|
|
super().__init__()
|
|
|
|
self.deconv = deconv
|
|
self.align_corners = align_corners
|
|
|
|
self.groups = 1
|
|
|
|
self.expand = expand
|
|
out_features = features
|
|
if self.expand is True:
|
|
out_features = features // 2
|
|
|
|
self.out_conv = nn.Conv2d(
|
|
features,
|
|
out_features,
|
|
kernel_size=1,
|
|
stride=1,
|
|
padding=0,
|
|
bias=True,
|
|
groups=1,
|
|
)
|
|
|
|
self.resConfUnit1 = ResidualConvUnit_custom(features, activation, bn)
|
|
self.resConfUnit2 = ResidualConvUnit_custom(features, activation, bn)
|
|
|
|
self.skip_add = nn.quantized.FloatFunctional()
|
|
|
|
def forward(self, *xs):
|
|
"""
|
|
Forward pass.
|
|
|
|
Returns:
|
|
tensor: output
|
|
"""
|
|
output = xs[0]
|
|
|
|
if len(xs) == 2:
|
|
res = self.resConfUnit1(xs[1])
|
|
output = self.skip_add.add(output, res)
|
|
# output += res
|
|
|
|
output = self.resConfUnit2(output)
|
|
|
|
output = nn.functional.interpolate(
|
|
output, scale_factor=2, mode="bilinear", align_corners=self.align_corners
|
|
)
|
|
|
|
output = self.out_conv(output)
|
|
|
|
return output
|