efficient_net

Complexity

Total Complexity

20

Size/Duplication

Total Lines	366
Duplicated Lines	0 %

Importance

Changes

0

8 Methods

Rating	Name	Duplication	Size	Complexity
A	EfficientNet.round_filters()	0	9	2
A	EfficientNet.build_efficient_net()	0	51	1
A	EfficientNet.round_repeats()	0	2	1
A	EfficientNet.correct_pad()	0	17	3
A	EfficientNet.call()	0	40	3
B	EfficientNet.block()	0	69	7
B	EfficientNet.__init__()	0	75	2
A	EfficientNet.get_config()	0	4	1

"""This script provides an example of using efficient for training."""

import os
import math
import numpy as np
import tensorflow as tf
# tf.config.experimental_run_functions_eagerly(True)

from tensorflow.keras import layers
from copy import deepcopy
from typing import List, Optional, Tuple, Union

from deepreg.model import layer
from deepreg.model.backbone import Backbone
from deepreg.model.backbone.local_net import LocalNet
from deepreg.model.backbone.u_net import UNet
from deepreg.model.layer import Extraction
from deepreg.registry import REGISTRY
from deepreg.train import train


EFFICIENTNET_PARAMS = {
    # model_name: (width_mult, depth_mult, image_size, dropout_rate, dropconnect_rate)
    "efficientnet-b0": (1.0, 1.0, 224, 0.2, 0.2),
    "efficientnet-b1": (1.0, 1.1, 240, 0.2, 0.2),
    "efficientnet-b2": (1.1, 1.2, 260, 0.3, 0.2),
    "efficientnet-b3": (1.2, 1.4, 300, 0.3, 0.2),
    "efficientnet-b4": (1.4, 1.8, 380, 0.4, 0.2),
    "efficientnet-b5": (1.6, 2.2, 456, 0.4, 0.2),
    "efficientnet-b6": (1.8, 2.6, 528, 0.5, 0.2),
    "efficientnet-b7": (2.0, 3.1, 600, 0.5, 0.2),
}

# Each Blocks Parameters
DEFAULT_BLOCKS_ARGS = [
    {'kernel_size': 3, 'repeats': 1, 'filters_in': 32, 'filters_out': 16,
     'expand_ratio': 1, 'id_skip': True, 'strides': 1, 'se_ratio': 0.25},
    {'kernel_size': 3, 'repeats': 2, 'filters_in': 16, 'filters_out': 24,
     'expand_ratio': 6, 'id_skip': True, 'strides': 2, 'se_ratio': 0.25},
    {'kernel_size': 5, 'repeats': 2, 'filters_in': 24, 'filters_out': 40,
     'expand_ratio': 6, 'id_skip': True, 'strides': 2, 'se_ratio': 0.25},
    {'kernel_size': 3, 'repeats': 3, 'filters_in': 40, 'filters_out': 80,
     'expand_ratio': 6, 'id_skip': True, 'strides': 2, 'se_ratio': 0.25},
    {'kernel_size': 5, 'repeats': 3, 'filters_in': 80, 'filters_out': 112,
     'expand_ratio': 6, 'id_skip': True, 'strides': 1, 'se_ratio': 0.25},
    {'kernel_size': 5, 'repeats': 4, 'filters_in': 112, 'filters_out': 192,
     'expand_ratio': 6, 'id_skip': True, 'strides': 2, 'se_ratio': 0.25},
    {'kernel_size': 3, 'repeats': 1, 'filters_in': 192, 'filters_out': 320,
     'expand_ratio': 6, 'id_skip': True, 'strides': 1, 'se_ratio': 0.25}
]

# Two Kernel Initializer
CONV_KERNEL_INITIALIZER = {
    'class_name': 'VarianceScaling',
    'config': {
        'scale': 2.0,
        'mode': 'fan_out',
        'distribution': 'normal'
    }
}

DENSE_KERNEL_INITIALIZER = {
    'class_name': 'VarianceScaling',
    'config': {
        'scale': 1. / 3.,
        'mode': 'fan_out',
        'distribution': 'uniform'
    }
}


@REGISTRY.register_backbone(name="efficient_net")
class EfficientNet(LocalNet):
    """
    Class that implements an Efficient-Net for image registration.

    Reference:
    - Author: Mingxing Tan, Quoc V. Le,
      EfficientNet: Rethinking Model Scaling for Convolutional Neural Networks
      https://arxiv.org/pdf/1905.11946.pdf
    """
    def __init__(
        self,
        image_size: tuple,
        num_channel_initial: int,
        extract_levels: Tuple[int, ...],
        out_kernel_initializer: str,
        out_activation: str,
        out_channels: int,
        depth: Optional[int] = None,
        pooling: bool = True,
        concat_skip: bool = False,
        width_coefficient: float = 1.0,
        depth_coefficient: float = 1.0,
        default_size: int = 224,
        dropout_rate: float = 0.2,
        drop_connect_rate: float = 0.2,
        depth_divisor: int = 8,
        name: str = "EfficientNet",
        **kwargs,
    ):
        """
        Init.

        Image is encoded gradually, i from level 0 to D,
        then it is decoded gradually, j from level D to 0.
        Some of the decoded levels are used for generating extractions.

        So, extract_levels are between [0, D].

        :param image_size: such as (dim1, dim2, dim3)
        :param num_channel_initial: number of initial channels.
        :param extract_levels: from which depths the output will be built.
        :param out_kernel_initializer: initializer to use for kernels.
        :param out_activation: activation to use at end layer.
        :param out_channels: number of channels for the extractions
        :param depth: depth of the encoder.
            If depth is not given, depth = max(extract_levels) will be used.
        :param pooling: for down-sampling, use non-parameterized
                        pooling if true, otherwise use conv3d
        :param concat_skip: when up-sampling, concatenate skipped
                            tensor if true, otherwise use addition
        :param width_coefficient: float, scaling coefficient for network width.
        :param depth_coefficient: float, scaling coefficient for network depth.
        :param default_size: int, default input image size.
        :param dropout_rate: float, dropout rate before final classifier layer.
        :param drop_connect_rate: float, dropout rate at skip connections.
        :param depth_divisor: int divisor for depth.
        :param name: name of the backbone.
        :param kwargs: additional arguments.
        """
        if depth is None:
            depth = max(extract_levels)
        kwargs["encode_kernel_sizes"] = [7] + [3] * depth
        super().__init__(
            image_size=image_size,
            num_channel_initial=num_channel_initial,
            depth=depth,
            extract_levels=extract_levels,
            out_kernel_initializer=out_kernel_initializer,
            out_activation=out_activation,
            out_channels=out_channels,
            use_additive_upsampling = False,
            pooling=pooling,
            concat_skip=concat_skip,
            name=name,
            **kwargs,
        )

        self.width_coefficient =  width_coefficient
        self.depth_coefficient = depth_coefficient
        self.default_size = default_size
        self.dropout_rate = dropout_rate
        self.drop_connect_rate = drop_connect_rate
        self.depth_divisor = depth_divisor
        self.activation_fn = tf.nn.swish

    def call(self, inputs: tf.Tensor, training=None, mask=None) -> tf.Tensor:
        """
        Build compute graph based on built layers.

        :param inputs: image batch, shape = (batch, f_dim1, f_dim2, f_dim3, ch)
        :param training: None or bool.
        :param mask: None or tf.Tensor.
        :return: shape = (batch, f_dim1, f_dim2, f_dim3, out_channels)
        """

        # encoding / down-sampling
        skips = []
        encoded = inputs
        print("before:", encoded.shape)

        for d in range(self._depth):
            skip = self._encode_convs[d](inputs=encoded, training=training)
            encoded = self._encode_pools[d](inputs=skip, training=training)
            skips.append(skip)
        print("skip: ", len(skip), self._depth)

The variable skip does not seem to be defined in case the for loop on line 173 is not entered. Are you sure this can never be the case?

        print("encode: ", encoded.shape)
        # bottom
        import efficientnet_3D.keras as efn 
        e_model = efn.EfficientNetB0(input_shape=(24, 24, 26, 64), weights='imagenet')
        decoded = e_model(encoded)
        # decoded = self.build_efficient_net(inputs=encoded)  # type: ignore

        print("decode: ", decoded.shape)
        # decoding / up-sampling
        outs = [decoded]
        for d in range(self._depth - 1, min(self._extract_levels) - 1, -1):
            decoded = self._decode_deconvs[d](inputs=decoded, training=training)
            decoded = self.build_skip_block()([decoded, skips[d]])
            decoded = self._decode_convs[d](inputs=decoded, training=training)
            outs = [decoded] + outs

        # output
        output = self._output_block(outs)  # type: ignore

        return output


    def build_efficient_net(self, inputs: tf.Tensor, training=None) -> tf.Tensor:
        """
        Builds graph based on built layers.

        :param inputs: shape = (batch, f_dim1, f_dim2, f_dim3, in_channels)
        :param training:
        :param mask:
        :return: shape = (batch, f_dim1, f_dim2, f_dim3, out_channels)
        """
        img_input = layers.Input(tensor=inputs, shape=self.image_size)
        bn_axis = 4  
        x = img_input
        # x = layers.ZeroPadding3D(padding=self.correct_pad(x, 3),
        #                         name='stem_conv_pad')(x)

        x = layers.Conv3D(self.round_filters(32), 3,
                        strides=1,
                        padding='same',
                        use_bias=False,
                        # kernel_initializer=CONV_KERNEL_INITIALIZER,
                        name='stem_conv')(x)
        x = layers.BatchNormalization(axis=bn_axis, name='stem_bn')(x)
        x = layers.Activation(self.activation_fn, name='stem_activation')(x)
        blocks_args = deepcopy(DEFAULT_BLOCKS_ARGS)

        b = 0
        # Calculate the number of blocks
        # blocks = float(sum(args['repeats'] for args in blocks_args))
        # for (i, args) in enumerate(blocks_args):
        #     assert args['repeats'] > 0
        #     args['filters_in'] = self.round_filters(args['filters_in'])
        #     args['filters_out'] = self.round_filters(args['filters_out'])

        #     for j in range(self.round_repeats(args.pop('repeats'))):
        #         if j > 0:
        #             args['strides'] = 1
        #             args['filters_in'] = args['filters_out']
        #         x = self.block(x, self.activation_fn, self.drop_connect_rate * b / blocks,
        #                 name='block{}{}_'.format(i + 1, chr(j + 97)), **args)
        #         b += 1
        
        x = layers.Conv3D(self.round_filters(128), 1,
                        padding='same',
                        use_bias=False,
                        # kernel_initializer=CONV_KERNEL_INITIALIZER,
                        name='top_conv')(x)
        x = layers.BatchNormalization(axis=bn_axis, name='top_bn')(x)
        x = layers.Activation(self.activation_fn, name='top_activation')(x)

        print("input.shape", inputs.shape, x.shape)
        return x

    def round_filters(self, filters):
        """Round number of filters based on depth multiplier."""
        filters *= self.width_coefficient
        divisor = self.depth_divisor
        new_filters = max(divisor, int(filters + divisor / 2) // divisor * divisor)
        # Make sure that round down does not go down by more than 10%.
        if new_filters < 0.9 * filters:
            new_filters += divisor
        return int(new_filters)

    def round_repeats(self, repeats):
        return int(math.ceil(self.depth_coefficient * repeats))

    def correct_pad(self, inputs, kernel_size):
        img_dim = 1
        input_size = backend.int_shape(inputs)[img_dim:(img_dim + 3)]

The variable backend does not seem to be defined.

        if isinstance(kernel_size, int):
            kernel_size = (kernel_size, kernel_size, kernel_size)

        if input_size[0] is None:
            adjust = (1, 1, 1)
        else:
            adjust = (1 - input_size[0] % 2, 1 - input_size[1] % 2, 1 - input_size[2] % 2)

        correct = (kernel_size[0] // 2, kernel_size[1] // 2, kernel_size[2] // 2)

        return ((correct[0] - adjust[0], correct[0]),
                (correct[1] - adjust[1], correct[1]),
                (correct[2] - adjust[2], correct[2]))

    def block(self, inputs, activation_fn=tf.nn.swish, drop_rate=0., name='',
            filters_in=32, filters_out=16, kernel_size=3, strides=1,
            expand_ratio=1, se_ratio=0., id_skip=True):

        bn_axis = 4

        filters = filters_in * expand_ratio

        # Inverted residuals
        if expand_ratio != 1:
            x = layers.Conv3D(filters, 1,
                            padding='same',
                            use_bias=False,
                            # kernel_initializer=CONV_KERNEL_INITIALIZER,
                            name=name + 'expand_conv')(inputs)
            x = layers.BatchNormalization(axis=bn_axis, name=name + 'expand_bn')(x)
            x = layers.Activation(activation_fn, name=name + 'expand_activation')(x)
        else:
            x = inputs

        # padding
        # if strides == 2:
        #     x = layers.ZeroPadding3D(padding=self.correct_pad(x, kernel_size),
        #                             name=name + 'dwconv_pad')(x)
        #     conv_pad = 'valid'
        # else:
        #     conv_pad = 'same'

        # TODO(Sicong): Find DepthwiseConv3D
        # x = layers.DepthwiseConv2D(kernel_size,
        #                         strides=strides,
        #                         padding=conv_pad,
        #                         use_bias=False,
        #                         depthwise_initializer=CONV_KERNEL_INITIALIZER,
        #                         name=name + 'dwconv')(x)
        x = layers.BatchNormalization(axis=bn_axis, name=name + 'bn')(x)
        x = layers.Activation(activation_fn, name=name + 'activation')(x)

        if 0 < se_ratio <= 1:
            filters_se = max(1, int(filters_in * se_ratio))
            se = layers.GlobalAveragePooling3D(name=name + 'se_squeeze')(x)
            se = layers.Reshape((1, 1, 1, filters), name=name + 'se_reshape')(se)
            se = layers.Conv3D(filters_se, 1,
                            padding='same',
                            activation=activation_fn,
                            # kernel_initializer=CONV_KERNEL_INITIALIZER,
                            name=name + 'se_reduce')(se)
            se = layers.Conv3D(filters, 1,
                            padding='same',
                            activation='sigmoid',
                            # kernel_initializer=CONV_KERNEL_INITIALIZER,
                            name=name + 'se_expand')(se)
            x = layers.multiply([x, se], name=name + 'se_excite')

        x = layers.Conv3D(filters_out, 1,
                        padding='same',
                        use_bias=False,
                        # kernel_initializer=CONV_KERNEL_INITIALIZER,
                        name=name + 'project_conv')(x)
        x = layers.BatchNormalization(axis=bn_axis, name=name + 'project_bn')(x)

        if (id_skip is True and strides == 1 and filters_in == filters_out):
            if drop_rate > 0:
                x = layers.Dropout(drop_rate,
                                noise_shape=None,
                                name=name + 'drop')(x)
            x = layers.add([x, inputs], name=name + 'add')

        return x

    def get_config(self) -> dict:
        """Return the config dictionary for recreating this class."""
        config = super().get_config()
        return config


if __name__ == "__main__":
    config_path = "examples/config_efficient_net.yaml"
    train(
        gpu="0",
        config_path=config_path,
        gpu_allow_growth=True,
        ckpt_path="",
    )