| import torch |
| import torch.nn as nn |
|
|
| def _make_encoder(backbone, features, use_pretrained, groups=1, expand=False, exportable=True): |
| if backbone == "efficientnet_lite3": |
| pretrained = _make_pretrained_efficientnet_lite3(use_pretrained, exportable=exportable) |
| scratch = _make_scratch([32, 48, 136, 384], features, groups=groups, expand=expand) |
| else: |
| print(f"Backbone '{backbone}' not implemented") |
| assert False |
| |
| 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==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, |
| trust_repo=True, |
| ) |
| 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 |
|
|
|
|
| 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==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==True: |
| out = self.bn1(out) |
| |
| out = self.activation(out) |
| out = self.conv2(out) |
| if self.bn==True: |
| out = self.bn2(out) |
|
|
| if self.groups > 1: |
| out = self.conv_merge(out) |
|
|
| return self.skip_add.add(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(FeatureFusionBlock_custom, self).__init__() |
|
|
| self.deconv = deconv |
| self.align_corners = align_corners |
|
|
| self.groups=1 |
|
|
| self.expand = expand |
| out_features = features |
| if self.expand==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 = 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 |
|
|
|
|
| class OutputConv(nn.Module): |
| """Output conv block. |
| """ |
|
|
| def __init__(self, features, groups, activation, non_negative): |
|
|
| super(OutputConv, self).__init__() |
|
|
| self.output_conv = nn.Sequential( |
| nn.Conv2d(features, features//2, kernel_size=3, stride=1, padding=1, groups=groups), |
| nn.Upsample(scale_factor=2, mode="bilinear"), |
| nn.Conv2d(features//2, 32, kernel_size=3, stride=1, padding=1), |
| activation, |
| nn.Conv2d(32, 1, kernel_size=1, stride=1, padding=0), |
| nn.ReLU(True) if non_negative else nn.Identity(), |
| nn.Identity(), |
| ) |
|
|
| def forward(self, x): |
| return self.output_conv(x) |
|
|
