Skip to content

instance_norm

InstanceNormalization

Bases: Layer

A layer for performing instance normalization.

This class is intentionally not @traceable (models and layers are handled by a different process).

This layer assumes that you are using the a tensor shaped like (Batch, Height, Width, Channels). See https://arxiv.org/abs/1607.08022 for details about this layer. The implementation here is borrowed from https://github.com/tensorflow/examples/blob/master/tensorflow_examples/models/pix2pix/pix2pix.py.

n = tfp.distributions.Normal(loc=10, scale=2)
x = n.sample(sample_shape=(1, 100, 100, 1))  # mean ~= 10, stddev ~= 2
m = fe.layers.tensorflow.InstanceNormalization()
y = m(x)  # mean ~= 0, stddev ~= 0

Parameters:

Name Type Description Default
epsilon float

A numerical stability constant added to the variance.

 1e-05
Source code in fastestimator/fastestimator/layers/tensorflow/instance_norm.py 
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
class InstanceNormalization(layers.Layer):
    """A layer for performing instance normalization.

    This class is intentionally not @traceable (models and layers are handled by a different process).

    This layer assumes that you are using the a tensor shaped like (Batch, Height, Width, Channels). See
    https://arxiv.org/abs/1607.08022 for details about this layer. The implementation here is borrowed from
    https://github.com/tensorflow/examples/blob/master/tensorflow_examples/models/pix2pix/pix2pix.py.

    ```python
    n = tfp.distributions.Normal(loc=10, scale=2)
    x = n.sample(sample_shape=(1, 100, 100, 1))  # mean ~= 10, stddev ~= 2
    m = fe.layers.tensorflow.InstanceNormalization()
    y = m(x)  # mean ~= 0, stddev ~= 0
    ```

    Args:
        epsilon: A numerical stability constant added to the variance.
    """
    def __init__(self, epsilon: float = 1e-5) -> None:
        super().__init__()
        self.epsilon = epsilon
        self.scale = None
        self.offset = None

    def get_config(self) -> Dict[str, Any]:
        return {'epsilon': self.epsilon}

    def build(self, input_shape: Tuple[int, int, int, int]) -> None:
        self.scale = self.add_weight(name='scale',
                                     shape=input_shape[-1:],
                                     initializer=tf.random_normal_initializer(0., 0.02),
                                     trainable=True)

        self.offset = self.add_weight(name='offset', shape=input_shape[-1:], initializer='zeros', trainable=True)

    def call(self, x: tf.Tensor) -> tf.Tensor:
        mean, variance = tf.nn.moments(x, axes=[1, 2], keepdims=True)
        inv = tf.math.rsqrt(variance + self.epsilon)
        normalized = (x - mean) * inv
        return self.scale * normalized + self.offset