clean/video/fakestormer/models/networks/pose_efficientNet.py

# -*- coding: utf-8 -*-
import math
import os
import sys

if not os.getcwd() in sys.path:
    sys.path.append(os.getcwd())

import torch
from torch import nn
from torch.nn import functional as F
from torch.utils import model_zoo

from ..builder import MODELS, build_model
from .backbones.efficientNet import (
    MemoryEfficientSwish,
    Swish,
    calculate_output_image_size,
    drop_connect,
    efficientnet_params,
    get_model_params,
    get_same_padding_conv2d,
    load_pretrained_weights,
    round_filters,
    round_repeats,
    url_map,
    url_map_advprop,
)
from .common import (
    BN_MOMENTUM,
    InceptionBlock,
    SELayer,
    Texture_Enhance,
    conv_block,
    point_wise_block,
)

VALID_MODELS = (
    "efficientnet-b0",
    "efficientnet-b1",
    "efficientnet-b2",
    "efficientnet-b3",
    "efficientnet-b4",
    "efficientnet-b5",
    "efficientnet-b6",
    "efficientnet-b7",
    "efficientnet-b8",
    # Support the construction of 'efficientnet-l2' without pretrained weights
    "efficientnet-l2",
)


class MBConvBlock(nn.Module):
    """Mobile Inverted Residual Bottleneck Block.
    Args:
        block_args (namedtuple): BlockArgs, defined in utils.py.
        global_params (namedtuple): GlobalParam, defined in utils.py.
        image_size (tuple or list): [image_height, image_width].
    References:
        [1] https://arxiv.org/abs/1704.04861 (MobileNet v1)
        [2] https://arxiv.org/abs/1801.04381 (MobileNet v2)
        [3] https://arxiv.org/abs/1905.02244 (MobileNet v3)
    """

    def __init__(self, block_args, global_params, image_size=None):
        super().__init__()
        self._block_args = block_args
        self._bn_mom = (
            1 - global_params.batch_norm_momentum
        )  # pytorch's difference from tensorflow
        self._bn_eps = global_params.batch_norm_epsilon
        self.has_se = (self._block_args.se_ratio is not None) and (
            0 < self._block_args.se_ratio <= 1
        )
        self.id_skip = (
            block_args.id_skip
        )  # whether to use skip connection and drop connect

        # Expansion phase (Inverted Bottleneck)
        inp = self._block_args.input_filters  # number of input channels
        oup = (
            self._block_args.input_filters * self._block_args.expand_ratio
        )  # number of output channels
        if self._block_args.expand_ratio != 1:
            Conv2d = get_same_padding_conv2d(image_size=image_size)
            self._expand_conv = Conv2d(
                in_channels=inp, out_channels=oup, kernel_size=1, bias=False
            )
            self._bn0 = nn.BatchNorm2d(
                num_features=oup, momentum=self._bn_mom, eps=self._bn_eps
            )
            # image_size = calculate_output_image_size(image_size, 1) <-- this wouldn't modify image_size

        # Depthwise convolution phase
        k = self._block_args.kernel_size
        s = self._block_args.stride
        Conv2d = get_same_padding_conv2d(image_size=image_size)
        self._depthwise_conv = Conv2d(
            in_channels=oup,
            out_channels=oup,
            groups=oup,  # groups makes it depthwise
            kernel_size=k,
            stride=s,
            bias=False,
        )
        self._bn1 = nn.BatchNorm2d(
            num_features=oup, momentum=self._bn_mom, eps=self._bn_eps
        )
        image_size = calculate_output_image_size(image_size, s)

        # Squeeze and Excitation layer, if desired
        if self.has_se:
            Conv2d = get_same_padding_conv2d(image_size=(1, 1))
            num_squeezed_channels = max(
                1, int(self._block_args.input_filters * self._block_args.se_ratio)
            )
            self._se_reduce = Conv2d(
                in_channels=oup, out_channels=num_squeezed_channels, kernel_size=1
            )
            self._se_expand = Conv2d(
                in_channels=num_squeezed_channels, out_channels=oup, kernel_size=1
            )

        # Pointwise convolution phase
        final_oup = self._block_args.output_filters
        Conv2d = get_same_padding_conv2d(image_size=image_size)
        self._project_conv = Conv2d(
            in_channels=oup, out_channels=final_oup, kernel_size=1, bias=False
        )
        self._bn2 = nn.BatchNorm2d(
            num_features=final_oup, momentum=self._bn_mom, eps=self._bn_eps
        )
        self._swish = MemoryEfficientSwish()

    def forward(self, inputs, drop_connect_rate=None):
        """MBConvBlock's forward function.
        Args:
            inputs (tensor): Input tensor.
            drop_connect_rate (bool): Drop connect rate (float, between 0 and 1).
        Returns:
            Output of this block after processing.
        """

        # Expansion and Depthwise Convolution
        x = inputs
        if self._block_args.expand_ratio != 1:
            x = self._expand_conv(inputs)
            x = self._bn0(x)
            x = self._swish(x)

        x = self._depthwise_conv(x)
        x = self._bn1(x)
        x = self._swish(x)

        # Squeeze and Excitation
        if self.has_se:
            x_squeezed = F.adaptive_avg_pool2d(x, 1)
            x_squeezed = self._se_reduce(x_squeezed)
            x_squeezed = self._swish(x_squeezed)
            x_squeezed = self._se_expand(x_squeezed)
            x = torch.sigmoid(x_squeezed) * x

        # Pointwise Convolution
        x = self._project_conv(x)
        x = self._bn2(x)

        # Skip connection and drop connect
        input_filters, output_filters = (
            self._block_args.input_filters,
            self._block_args.output_filters,
        )
        if (
            self.id_skip
            and self._block_args.stride == 1
            and input_filters == output_filters
        ):
            # The combination of skip connection and drop connect brings about stochastic depth.
            if drop_connect_rate:
                x = drop_connect(x, p=drop_connect_rate, training=self.training)
            x = x + inputs  # skip connection
        return x

    def set_swish(self, memory_efficient=True):
        """Sets swish function as memory efficient (for training) or standard (for export).
        Args:
            memory_efficient (bool): Whether to use memory-efficient version of swish.
        """
        self._swish = MemoryEfficientSwish() if memory_efficient else Swish()


@MODELS.register_module()
class EfficientNet(nn.Module):
    """EfficientNet model.
       Most easily loaded with the .from_name or .from_pretrained methods.
    Args:
        blocks_args (list[namedtuple]): A list of BlockArgs to construct blocks.
        global_params (namedtuple): A set of GlobalParams shared between blocks.
    References:
        [1] https://arxiv.org/abs/1905.11946 (EfficientNet)
    Example:
        >>> import torch
        >>> from efficientnet.model import EfficientNet
        >>> inputs = torch.rand(1, 3, 224, 224)
        >>> model = EfficientNet.from_pretrained('efficientnet-b0')
        >>> model.eval()
        >>> outputs = model(inputs)
    """

    def __init__(self, blocks_args=None, global_params=None):
        super().__init__()
        assert isinstance(blocks_args, list), "blocks_args should be a list"
        assert len(blocks_args) > 0, "block args must be greater than 0"
        self._global_params = global_params
        self._blocks_args = blocks_args

        # Batch norm parameters
        bn_mom = 1 - self._global_params.batch_norm_momentum
        bn_eps = self._global_params.batch_norm_epsilon

        # Get stem static or dynamic convolution depending on image size
        image_size = global_params.image_size
        Conv2d = get_same_padding_conv2d(image_size=image_size)

        # Stem
        in_channels = 3  # rgb
        out_channels = round_filters(
            32, self._global_params
        )  # number of output channels
        self._conv_stem = Conv2d(
            in_channels, out_channels, kernel_size=3, stride=2, bias=False
        )
        self._bn0 = nn.BatchNorm2d(
            num_features=out_channels, momentum=bn_mom, eps=bn_eps
        )
        image_size = calculate_output_image_size(image_size, 2)

        # Build blocks
        self._blocks = nn.ModuleList([])
        for block_args in self._blocks_args:

            # Update block input and output filters based on depth multiplier.
            block_args = block_args._replace(
                input_filters=round_filters(
                    block_args.input_filters, self._global_params
                ),
                output_filters=round_filters(
                    block_args.output_filters, self._global_params
                ),
                num_repeat=round_repeats(block_args.num_repeat, self._global_params),
            )

            # The first block needs to take care of stride and filter size increase.
            self._blocks.append(
                MBConvBlock(block_args, self._global_params, image_size=image_size)
            )
            image_size = calculate_output_image_size(image_size, block_args.stride)
            if block_args.num_repeat > 1:  # modify block_args to keep same output size
                block_args = block_args._replace(
                    input_filters=block_args.output_filters, stride=1
                )
            for _ in range(block_args.num_repeat - 1):
                self._blocks.append(
                    MBConvBlock(block_args, self._global_params, image_size=image_size)
                )
                # image_size = calculate_output_image_size(image_size, block_args.stride)  # stride = 1

        # Head
        in_channels = block_args.output_filters  # output of final block
        out_channels = round_filters(1280, self._global_params)
        Conv2d = get_same_padding_conv2d(image_size=image_size)
        self._conv_head = Conv2d(in_channels, out_channels, kernel_size=1, bias=False)
        self._bn1 = nn.BatchNorm2d(
            num_features=out_channels, momentum=bn_mom, eps=bn_eps
        )

        # Final linear layer
        self._avg_pooling = nn.AdaptiveAvgPool2d(1)
        if self._global_params.include_top:
            self._dropout = nn.Dropout(self._global_params.dropout_rate)
            self._fc = nn.Linear(out_channels, self._global_params.num_classes)

        # Heatmap Decoder Construction
        if self._global_params.include_hm_decoder:
            print("Constructing the heatmap Decoder!")
            self.efpn = self._global_params.efpn
            self.tfpn = self._global_params.tfpn

            assert not (
                self.efpn and self.tfpn
            ), "Only one of E-FPN or FPN is intergrated!"

            self.se_layer = self._global_params.se_layer
            # self.hm_decoder_filters = [1792, 448, 160, 56] if self.fpn else [1792, 256, 256, 128]
            self.hm_decoder_filters = [1792, 448, 160, 56]
            num_kernels = [4, 4, 4, 4] if (self.efpn or self.tfpn) else [4, 4, 4]
            self._dropout = nn.Dropout(self._global_params.dropout_rate)
            self._sigmoid = nn.Sigmoid()
            self._relu = nn.ReLU(inplace=True)
            self._relu1 = nn.ReLU(inplace=False)
            self.deconv_with_bias = False

            # if self._global_params.use_c2:
            #     self.inception_block_4 = InceptionBlock(224, 224, stride=1, pool_size=3)
            if self._global_params.use_c3:
                self.inception_block = InceptionBlock(112, 112, stride=1, pool_size=3)
            else:
                self.inception_block = InceptionBlock(56, 56, stride=1, pool_size=3)

            self.heads = self._global_params.heads
            n_deconv = len(self.hm_decoder_filters)
            self.fpn_layers = [
                self._global_params.use_c51,
                self._global_params.use_c4,
                self._global_params.use_c3,
            ]

            if self.efpn or self.tfpn:
                for idx in range(n_deconv):
                    in_decod_filters = self.hm_decoder_filters[idx]

                    if idx == 0:
                        out_decod_filters = self.hm_decoder_filters[idx + 1]
                        deconv = nn.Sequential(
                            conv_block(
                                in_decod_filters,
                                out_decod_filters,
                                (3, 3),
                                stride=1,
                                padding=1,
                            ),
                        )
                    else:
                        in_decod_filters = (
                            in_decod_filters * 2
                            if self.fpn_layers[idx - 1]
                            else in_decod_filters
                        )
                        kernel, padding, output_padding = self._get_deconv_cfg(
                            num_kernels[idx]
                        )

                        if idx + 1 < n_deconv:
                            out_decod_filters = self.hm_decoder_filters[idx + 1]
                            deconv = nn.Sequential(
                                conv_block(
                                    in_decod_filters,
                                    out_decod_filters,
                                    (3, 3),
                                    stride=1,
                                    padding=1,
                                ),
                                nn.ConvTranspose2d(
                                    in_channels=out_decod_filters,
                                    out_channels=out_decod_filters,
                                    kernel_size=kernel,
                                    stride=2,
                                    padding=padding,
                                    output_padding=output_padding,
                                    bias=self.deconv_with_bias,
                                ),
                                nn.BatchNorm2d(out_decod_filters, momentum=BN_MOMENTUM),
                            )
                        else:
                            out_decod_filters = in_decod_filters
                            deconv = nn.Sequential(
                                self.inception_block,
                                # conv_block(in_decod_filters, out_decod_filters, (3,3), stride=1, padding=1),
                                nn.ConvTranspose2d(
                                    in_channels=out_decod_filters,
                                    out_channels=out_decod_filters,
                                    kernel_size=kernel,
                                    stride=2,
                                    padding=padding,
                                    output_padding=output_padding,
                                    bias=self.deconv_with_bias,
                                ),
                                nn.BatchNorm2d(out_decod_filters, momentum=BN_MOMENTUM),
                            )

                            # In case of using C2, this conv to apply to C2 features to get the same filters of the last deconv
                            if self._global_params.use_c2:
                                if self._global_params.norm_c2:
                                    self.texture_enhance = Texture_Enhance(32)
                                    # self.conv_c2 = point_wise_block(128, out_decod_filters)
                                    self.conv_c2 = conv_block(
                                        128,
                                        out_decod_filters,
                                        (3, 3),
                                        stride=1,
                                        padding=1,
                                    )
                                else:
                                    self.conv_c2 = conv_block(
                                        32,
                                        out_decod_filters,
                                        (3, 3),
                                        stride=1,
                                        padding=1,
                                    )
                    if self.se_layer:
                        se = SELayer(channel=out_decod_filters * 2)
                        self.__setattr__(f"se_layer_{idx+1}", se)

                    self.__setattr__(f"deconv_{idx+1}", deconv)
            else:
                self.deconv_layers = self._make_deconv_layer(
                    len(num_kernels),
                    self.hm_decoder_filters,
                    num_kernels,
                )

            for head, num_output in self.heads.items():
                head_conv = int(self._global_params.head_conv)
                num_output = int(num_output)
                if self._global_params.use_c2:
                    assert (
                        self._global_params.efpn or self._global_params.tfpn
                    ), "FPN Design must be set active!"
                    assert (
                        self._global_params.use_c3
                    ), "C3 must be utilized for FPN intergration of C2"
                    in_head_filters = self.hm_decoder_filters[-1] * 4
                elif self._global_params.use_c3:
                    in_head_filters = self.hm_decoder_filters[-1] * 2
                else:
                    in_head_filters = self.hm_decoder_filters[-1]

                if head_conv > 0:
                    if head != "cls":
                        fc = nn.Sequential(
                            nn.Conv2d(
                                in_head_filters,
                                head_conv,
                                kernel_size=3,
                                padding=1,
                                bias=True,
                            ),
                            nn.BatchNorm2d(head_conv),
                            nn.ReLU(inplace=True),
                            nn.Conv2d(
                                head_conv,
                                num_output,
                                kernel_size=1,
                                stride=1,
                                padding=0,
                            ),
                        )
                    else:
                        fc = nn.Sequential(
                            nn.Conv2d(
                                in_head_filters,
                                head_conv,
                                kernel_size=3,
                                padding=1,
                                bias=True,
                            ),
                            nn.BatchNorm2d(head_conv, momentum=BN_MOMENTUM),
                            nn.ReLU(inplace=True),
                            # nn.Conv2d(head_conv, num_output, kernel_size=1,
                            #             stride=1, padding=0, bias=True),
                            # nn.BatchNorm2d(num_output),
                            # nn.ReLU(inplace=True),
                            # nn.AdaptiveMaxPool2d(head_conv//4),
                            nn.AdaptiveAvgPool2d(1),
                            nn.Flatten(),
                            # nn.Linear((head_conv//4)**2, head_conv, bias=True),
                            # nn.BatchNorm1d(head_conv, momentum=BN_MOMENTUM),
                            # nn.ReLU(inplace=True),
                            nn.Linear(head_conv, num_output, bias=True),
                            # nn.Sigmoid(),
                            # nn.Softmax(dim=-1)
                        )
                else:
                    fc = nn.Conv2d(
                        in_channels=in_head_filters,
                        out_channels=num_output,
                        kernel_size=1,
                        stride=1,
                        padding=0,
                    )
                self.__setattr__(head, fc)

        # set activation to memory efficient swish by default
        self._swish = MemoryEfficientSwish()

    def _get_deconv_cfg(self, deconv_kernel):
        if deconv_kernel == 4:
            padding = 1
            output_padding = 0
        elif deconv_kernel == 3:
            padding = 1
            output_padding = 1
        elif deconv_kernel == 2:
            padding = 0
            output_padding = 0

        return deconv_kernel, padding, output_padding

    def _make_deconv_layer(self, num_layers, num_filters, num_kernels):
        assert num_layers == (
            len(num_filters) - 1
        ), "ERROR: num_deconv_layers is different len(num_deconv_filters)"
        assert num_layers == len(
            num_kernels
        ), "ERROR: num_deconv_layers is different len(num_deconv_filters)"

        layers = []
        for i in range(num_layers):
            kernel, padding, output_padding = self._get_deconv_cfg(num_kernels[i])

            in_planes = num_filters[i]
            out_planes = num_filters[i + 1]

            layers.append(
                nn.Sequential(
                    nn.ConvTranspose2d(
                        in_channels=in_planes,
                        out_channels=out_planes,
                        kernel_size=kernel,
                        stride=2,
                        padding=padding,
                        output_padding=output_padding,
                        bias=self.deconv_with_bias,
                    ),
                    nn.BatchNorm2d(out_planes, momentum=BN_MOMENTUM),
                    nn.ReLU(inplace=True),
                )
            )

        return nn.Sequential(*layers)

    def set_swish(self, memory_efficient=True):
        """Sets swish function as memory efficient (for training) or standard (for export).
        Args:
            memory_efficient (bool): Whether to use memory-efficient version of swish.
        """
        self._swish = MemoryEfficientSwish() if memory_efficient else Swish()
        for block in self._blocks:
            block.set_swish(memory_efficient)

    def extract_endpoints(self, inputs):
        """Use convolution layer to extract features
        from reduction levels i in [1, 2, 3, 4, 5].
        Args:
            inputs (tensor): Input tensor.
        Returns:
            Dictionary of last intermediate features
            with reduction levels i in [1, 2, 3, 4, 5].
            Example:
                >>> import torch
                >>> from efficientnet.model import EfficientNet
                >>> inputs = torch.rand(1, 3, 224, 224)
                >>> model = EfficientNet.from_pretrained('efficientnet-b0')
                >>> endpoints = model.extract_endpoints(inputs)
                >>> print(endpoints['reduction_1'].shape)  # torch.Size([1, 16, 112, 112])
                >>> print(endpoints['reduction_2'].shape)  # torch.Size([1, 24, 56, 56])
                >>> print(endpoints['reduction_3'].shape)  # torch.Size([1, 40, 28, 28])
                >>> print(endpoints['reduction_4'].shape)  # torch.Size([1, 112, 14, 14])
                >>> print(endpoints['reduction_5'].shape)  # torch.Size([1, 320, 7, 7])
                >>> print(endpoints['reduction_6'].shape)  # torch.Size([1, 1280, 7, 7])
        """
        endpoints = dict()

        # Stem
        x = self._swish(self._bn0(self._conv_stem(inputs)))
        prev_x = x

        # Blocks
        for idx, block in enumerate(self._blocks):
            drop_connect_rate = self._global_params.drop_connect_rate
            if drop_connect_rate:
                drop_connect_rate *= float(idx) / len(
                    self._blocks
                )  # scale drop connect_rate
            x = block(x, drop_connect_rate=drop_connect_rate)
            # print('Prev', prev_x.size())
            # print('X', x.size())
            if prev_x.size(2) > x.size(2):
                endpoints["reduction_{}".format(len(endpoints) + 1)] = prev_x
            elif idx == len(self._blocks) - 1:
                endpoints["reduction_{}".format(len(endpoints) + 1)] = x
            prev_x = x

        # Head
        x = self._swish(self._bn1(self._conv_head(x)))
        endpoints["reduction_{}".format(len(endpoints) + 1)] = x

        return endpoints

    def extract_features(self, inputs):
        """use convolution layer to extract feature .
        Args:
            inputs (tensor): Input tensor.
        Returns:
            Output of the final convolution
            layer in the efficientnet model.
        """
        # Stem
        x = self._swish(self._bn0(self._conv_stem(inputs)))

        # Blocks
        for idx, block in enumerate(self._blocks):
            drop_connect_rate = self._global_params.drop_connect_rate
            if drop_connect_rate:
                drop_connect_rate *= float(idx) / len(
                    self._blocks
                )  # scale drop connect_rate
            x = block(x, drop_connect_rate=drop_connect_rate)

        # Head
        x = self._swish(self._bn1(self._conv_head(x)))

        return x

    def forward(self, inputs):
        """EfficientNet's forward function.
           Calls extract_features to extract features, applies final linear layer, and returns logits.
        Args:
            inputs (tensor): Input tensor.
        Returns:
            Output of this model after processing.
        """
        # Convolution layers
        # x = self.extract_features(inputs)
        endpoints = self.extract_endpoints(inputs)
        x1 = endpoints["reduction_6"]
        x2 = endpoints["reduction_5"]
        x3 = endpoints["reduction_4"]
        x4 = endpoints["reduction_3"]
        x5 = endpoints["reduction_2"]
        x = x1

        if self._global_params.include_top:
            # Pooling and final linear layer
            x = self._avg_pooling(x)

            x = x.flatten(start_dim=1)
            x = self._dropout(x)
            x = self._fc(x)
            res = {}
            res["cls"] = x

            return res

        if self._global_params.include_hm_decoder:
            x1 = self._dropout(x1)
            x2 = self._dropout(x2)
            x3 = self._dropout(x3)
            x4 = self._dropout(x4)

            if self.efpn:
                assert (
                    self._global_params.use_c51
                ), "C51 must be utilized for FPN intergration"

                x = self.__getattr__("deconv_1")(x1)

                if self._global_params.use_c51:
                    x_weighted = self._sigmoid(x)
                    x_inv = torch.sub(1, x_weighted, alpha=1)
                    x2_ = torch.multiply(x_inv, x2)
                    x = torch.cat([x, x2_], dim=1)

                    if self.se_layer:
                        x = self.__getattr__("se_layer_1")(x)
                else:
                    x = self._relu(x)

                x = self.__getattr__("deconv_2")(x)

                if self._global_params.use_c4:
                    x_weighted = self._sigmoid(x)
                    x_inv = torch.sub(1, x_weighted, alpha=1)
                    x3_ = torch.multiply(x_inv, x3)
                    x = torch.cat([x, x3_], dim=1)

                    if self.se_layer:
                        x = self.__getattr__("se_layer_2")(x)
                else:
                    x = self._relu(x)

                x = self.__getattr__("deconv_3")(x)

                if self._global_params.use_c3:
                    assert (
                        self._global_params.use_c4
                    ), "C4 must be utilized for FPN intergration of C3"

                    x_weighted = self._sigmoid(x)
                    x_inv = torch.sub(1, x_weighted, alpha=1)
                    x4_ = torch.multiply(x_inv, x4)
                    x = torch.cat([x, x4_], dim=1)

                    if self.se_layer:
                        x = self.__getattr__("se_layer_3")(x)
                else:
                    x = self._relu(x)

                x = self.__getattr__("deconv_4")(x)

                if not self._global_params.use_c2:
                    x = self._relu(x)
                else:
                    assert (
                        self._global_params.use_c3
                    ), "C3 must be utilized for FPN intergration of C2"

                    x5 = self._dropout(x5)
                    if self._global_params.norm_c2:
                        x5_ = self.texture_enhance(x5, (16, 16))
                        x5_ = self.conv_c2(x5_)
                    else:
                        x5_ = self.conv_c2(x5)

                    x_weighted = self._sigmoid(x)
                    x_inv = torch.sub(1, x_weighted, alpha=1)
                    x5_ = torch.multiply(x_inv, x5_)
                    x = torch.cat([x, x5_], dim=1)

                    # # Adding multi receptive fields
                    # x = self.inception_block_4(x)

                    if self.se_layer:
                        x = self.__getattr__("se_layer_4")(x)
            elif self.tfpn:
                assert (
                    self._global_params.use_c51
                ), "C51 must be utilized for FPN intergration"
                x = self.__getattr__("deconv_1")(x1)
                x = self._relu1(x)
                x = torch.cat([x, x2], dim=1)

                x = self.__getattr__("deconv_2")(x)
                if not self._global_params.use_c4:
                    x = self._relu1(x)
                else:
                    x = torch.cat([x, x3], dim=1)

                x = self.__getattr__("deconv_3")(x)
                if not self._global_params.use_c3:
                    x = self._relu1(x)
                else:
                    assert (
                        self._global_params.use_c4
                    ), "C4 must be utilized for FPN intergration of C3"
                    x = torch.cat([x, x4], dim=1)

                x = self.__getattr__("deconv_4")(x)
                if not self._global_params.use_c2:
                    x = self._relu(x)
                else:
                    assert (
                        self._global_params.use_c3
                    ), "C3 must be utilized for FPN intergration of C2"
                    x5 = self._dropout(x5)
                    x5 = self.conv_c2(x5)
                    x = self._relu1(x)
                    x = torch.cat([x, x5], dim=1)
            else:
                x = self.deconv_layers(x1)

            ret = {}
            for head in self.heads:
                ret[head] = self.__getattr__(head)(x)

            return [ret]

    @classmethod
    def from_name(cls, model_name, in_channels=3, **override_params):
        """Create an efficientnet model according to name.
        Args:
            model_name (str): Name for efficientnet.
            in_channels (int): Input data's channel number.
            override_params (other key word params):
                Params to override model's global_params.
                Optional key:
                    'width_coefficient', 'depth_coefficient',
                    'image_size', 'dropout_rate',
                    'num_classes', 'batch_norm_momentum',
                    'batch_norm_epsilon', 'drop_connect_rate',
                    'depth_divisor', 'min_depth'
        Returns:
            An efficientnet model.
        """
        cls._check_model_name_is_valid(model_name)
        blocks_args, global_params = get_model_params(model_name, override_params)
        model = cls(blocks_args, global_params)
        model._change_in_channels(in_channels)
        return model

    @classmethod
    def from_pretrained(
        cls,
        model_name,
        weights_path=None,
        advprop=False,
        in_channels=3,
        num_classes=1000,
        **override_params,
    ):
        """Create an efficientnet model according to name.
        Args:
            model_name (str): Name for efficientnet.
            weights_path (None or str):
                str: path to pretrained weights file on the local disk.
                None: use pretrained weights downloaded from the Internet.
            advprop (bool):
                Whether to load pretrained weights
                trained with advprop (valid when weights_path is None).
            in_channels (int): Input data's channel number.
            num_classes (int):
                Number of categories for classification.
                It controls the output size for final linear layer.
            override_params (other key word params):
                Params to override model's global_params.
                Optional key:
                    'width_coefficient', 'depth_coefficient',
                    'image_size', 'dropout_rate',
                    'batch_norm_momentum',
                    'batch_norm_epsilon', 'drop_connect_rate',
                    'depth_divisor', 'min_depth'
        Returns:
            A pretrained efficientnet model.
        """
        model = cls.from_name(model_name, num_classes=num_classes, **override_params)
        load_pretrained_weights(
            model,
            model_name,
            weights_path=weights_path,
            load_fc=((num_classes == 1000) and (model._global_params.include_top)),
            advprop=advprop,
        )
        model._change_in_channels(in_channels)
        return model

    @classmethod
    def get_image_size(cls, model_name):
        """Get the input image size for a given efficientnet model.
        Args:
            model_name (str): Name for efficientnet.
        Returns:
            Input image size (resolution).
        """
        cls._check_model_name_is_valid(model_name)
        _, _, res, _ = efficientnet_params(model_name)
        return res

    @classmethod
    def _check_model_name_is_valid(cls, model_name):
        """Validates model name.
        Args:
            model_name (str): Name for efficientnet.
        Returns:
            bool: Is a valid name or not.
        """
        if model_name not in VALID_MODELS:
            raise ValueError("model_name should be one of: " + ", ".join(VALID_MODELS))

    def _change_in_channels(self, in_channels):
        """Adjust model's first convolution layer to in_channels, if in_channels not equals 3.
        Args:
            in_channels (int): Input data's channel number.
        """
        if in_channels != 3:
            Conv2d = get_same_padding_conv2d(image_size=self._global_params.image_size)
            out_channels = round_filters(32, self._global_params)
            self._conv_stem = Conv2d(
                in_channels, out_channels, kernel_size=3, stride=2, bias=False
            )


@MODELS.register_module()
class PoseEfficientNet(EfficientNet):
    def __init__(self, model_name, in_channels=3, **override_params):
        self.model_name = model_name
        self.in_channels = in_channels

        # Initialize Parent Class
        super()._check_model_name_is_valid(model_name)
        blocks_args, global_params = get_model_params(model_name, override_params)
        super().__init__(blocks_args, global_params)

    @classmethod
    def from_name(cls, model_name, in_channels, **override_params):
        return NotImplemented

    @classmethod
    def from_pretrained(
        cls,
        model_name,
        weights_path,
        advprop,
        in_channels,
        num_classes,
        **override_params,
    ):
        return NotImplemented

    def _change_in_channels(self, in_channels):
        return NotImplemented

    def init_weights(self, pretrained=False, advprop=False, verbose=True):
        if pretrained:
            url_map_ = url_map_advprop if advprop else url_map
            state_dict = model_zoo.load_url(url_map_[self.model_name])
            state_dict.pop("_fc.weight")
            state_dict.pop("_fc.bias")
            self.load_state_dict(state_dict, strict=False)

        # Initialize weights for Deconvolution Layer
        if self._global_params.include_hm_decoder:
            if self.efpn or self.tfpn:
                deconv_layers = [
                    self.deconv_1,
                    self.deconv_2,
                    self.deconv_3,
                    self.deconv_4,
                ]
            else:
                deconv_layers = self.deconv_layers

            for layer in deconv_layers:
                for _, m in layer.named_modules():
                    if isinstance(m, nn.ConvTranspose2d):
                        n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                        m.weight.data.normal_(0, math.sqrt(2.0 / n))
                        if self.deconv_with_bias:
                            nn.init.constant_(m.bias, 0)
                    elif isinstance(m, nn.BatchNorm2d):
                        nn.init.constant_(m.weight, 1)
                        nn.init.constant_(m.bias, 0)

            # Init head parameters
            for head in self.heads:
                final_layer = self.__getattr__(head)
                for i, m in enumerate(final_layer.modules()):
                    if isinstance(m, nn.Conv2d):
                        if m.weight.shape[0] == self.heads[head]:
                            if "hm" in head:
                                nn.init.constant_(m.bias, -2.19)
                            else:
                                # nn.init.normal_(m.weight, std=0.001)
                                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                                m.weight.data.normal_(0, math.sqrt(2.0 / n))
                                nn.init.constant_(m.bias, 0)

        self._change_in_channels(in_channels=self.in_channels)
        if verbose:
            print("Loaded pretrained weights for {}".format(self.model_name))


if __name__ == "__main__":
    cfg = dict(
        type="PoseEfficientNet",
        model_name="efficientnet-b4",
        include_top=False,
        include_hm_decoder=True,
        head_conv=64,
        heads={"hm": 1, "cls": 1, "cstency": 256},
        use_c2=True,
    )
    model = build_model(cfg, MODELS)
    model.init_weights(pretrained=True)
    model.eval()
    inputs = torch.rand((1, 3, 384, 384))

    for i, (n, p) in enumerate(model.named_parameters()):
        print(i, n)

    # To show the whole pose EFN model outputs shape
    x = model(inputs)[0]
    for head in x.keys():
        print(f"{head} shape is --- {x[head].shape}")

    # To show the endpoints features shape
    # endpoints = model.extract_endpoints(inputs)
    # for k in endpoints.keys():
    #     print(endpoints[k].shape)