faceswap/tests/lib/model/layers_test.py

#!/usr/bin/env python3
""" Tests for Faceswap Custom Layers.

Adapted from Keras tests.
"""


import pytest
import numpy as np

from numpy.testing import assert_allclose

from keras import device, Input, Model, backend as K

from lib.model import layers
from lib.utils import get_backend
from tests.utils import has_arg


# pylint:disable=dangerous-default-value,too-many-locals,too-many-branches
def layer_test(layer_cls,  # noqa:C901
               kwargs={},
               input_shape=None,
               input_dtype=None,
               input_data=None,
               expected_output=None,
               expected_output_dtype=None,
               fixed_batch_size=False):
    """Test routine for a layer with a single input tensor
    and single output tensor.
    """
    with device("cpu"):
        # generate input data
        # pylint:disable=duplicate-code
        if input_data is None:
            assert input_shape
            if not input_dtype:
                input_dtype = K.floatx()
            input_data_shape = list(input_shape)
            for i, var_e in enumerate(input_data_shape):
                if var_e is None:
                    input_data_shape[i] = np.random.randint(1, 4)
            input_data = 10 * np.random.random(input_data_shape)
            input_data = input_data.astype(input_dtype)
        else:
            if input_shape is None:
                input_shape = input_data.shape
            if input_dtype is None:
                input_dtype = input_data.dtype
        if expected_output_dtype is None:
            expected_output_dtype = input_dtype

        # instantiation
        layer = layer_cls(**kwargs)

        # test get_weights , set_weights at layer level
        weights = layer.get_weights()
        layer.set_weights(weights)

        layer.build(input_shape)
        expected_output_shape = layer.compute_output_shape(input_shape)

        # test in functional API
        if fixed_batch_size:
            inp = Input(batch_shape=input_shape, dtype=input_dtype)
        else:
            inp = Input(shape=input_shape[1:], dtype=input_dtype)
        outp = layer(inp)
        assert outp.dtype == expected_output_dtype

        # check with the functional API
        model = Model(inp, outp)

        actual_output = model.predict(input_data, verbose=0)  # type:ignore
        actual_output_shape = actual_output.shape
        for expected_dim, actual_dim in zip(expected_output_shape,
                                            actual_output_shape):
            if expected_dim is not None:
                assert expected_dim == actual_dim

        if expected_output is not None:
            assert_allclose(actual_output, expected_output, rtol=1e-3)

        # test serialization, weight setting at model level
        model_config = model.get_config()
        recovered_model = model.__class__.from_config(model_config)
        if model.weights:
            weights = model.get_weights()
            recovered_model.set_weights(weights)
            _output = recovered_model.predict(input_data, verbose=0)  # type:ignore
            assert_allclose(_output, actual_output, rtol=1e-3)

        # test training mode (e.g. useful when the layer has a
        # different behavior at training and testing time).
        if has_arg(layer.call, 'training'):
            model.compile('rmsprop', 'mse')
            model.train_on_batch(input_data, actual_output)

        # test instantiation from layer config
        layer_config = layer.get_config()
        layer = layer.__class__.from_config(layer_config)

        # for further checks in the caller function
        return actual_output


@pytest.mark.parametrize('dummy', [None], ids=[get_backend().upper()])
def test_global_min_pooling_2d(dummy):  # pylint:disable=unused-argument
    """ Global Min Pooling 2D layer test """
    layer_test(layers.GlobalMinPooling2D, input_shape=(2, 4, 4, 1024))


@pytest.mark.parametrize('dummy', [None], ids=[get_backend().upper()])
def test_global_std_pooling_2d(dummy):  # pylint:disable=unused-argument
    """ Global Standard Deviation Pooling 2D layer test """
    layer_test(layers.GlobalStdDevPooling2D, input_shape=(2, 4, 4, 1024))


@pytest.mark.parametrize('dummy', [None], ids=[get_backend().upper()])
def test_k_resize_images(dummy):  # pylint:disable=unused-argument
    """ Global Standard Deviation Pooling 2D layer test """
    layer_test(layers.KResizeImages, input_shape=(2, 4, 4, 1024))


@pytest.mark.parametrize('dummy', [None], ids=[get_backend().upper()])
def test_l2_normalize(dummy):  # pylint:disable=unused-argument
    """ L2 Normalize layer test """
    layer_test(layers.L2Normalize, kwargs={"axis": 1}, input_shape=(2, 4, 4, 1024))


@pytest.mark.parametrize('dummy', [None], ids=[get_backend().upper()])
def test_pixel_shuffler(dummy):  # pylint:disable=unused-argument
    """ Pixel Shuffler layer test """
    layer_test(layers.PixelShuffler, input_shape=(2, 4, 4, 1024))


@pytest.mark.parametrize('dummy', [None], ids=[get_backend().upper()])
def test_quick_gelu(dummy):  # pylint:disable=unused-argument
    """ Global Standard Deviation Pooling 2D layer test """
    layer_test(layers.QuickGELU, input_shape=(2, 4, 4, 1024))


@pytest.mark.parametrize('dummy', [None], ids=[get_backend().upper()])
def test_reflection_padding_2d(dummy):  # pylint:disable=unused-argument
    """ Reflection Padding 2D layer test """
    layer_test(layers.ReflectionPadding2D, input_shape=(2, 4, 4, 512))


@pytest.mark.parametrize('dummy', [None], ids=[get_backend().upper()])
def test_swish(dummy):  # pylint:disable=unused-argument
    """ Swish activation layer test """
    layer_test(layers.Swish, input_shape=(2, 4, 4, 1024))


_PARAMS = ("multiply", "truediv", "add", "subtract")
_IDS = [f"{x}[{get_backend().upper()}]" for x in _PARAMS]


@pytest.mark.parametrize("operation", _PARAMS, ids=_IDS)
def test_scalar_op(operation):
    """ Scalar operation layer test """
    val = 2.0
    np_ops = {"multiply": np.multiply,
              "truediv": np.true_divide,
              "add": np.add,
              "subtract": np.subtract}
    input_data = np.random.random((2, 4, 4, 1024)).astype("float32")
    output_data = np_ops[operation](input_data, val)
    layer_test(layers.ScalarOp,
               kwargs={"operation": operation, "value": val},
               input_data=input_data,
               expected_output=output_data)