deeplabcut.pose_estimation_tensorflow.nnets.utils

Classes:

Name	Description
BatchNormalization	Fixed default name of BatchNormalization to match TpuBatchNormalization.
DepthwiseConv2D	Wrap keras DepthwiseConv2D to tf.layers.
TpuBatchNormalization	Cross replica batch normalization.

Functions:

Name	Description
build_learning_rate	Build learning rate.
build_optimizer	Build optimizer.
drop_connect	Apply drop connect.

BatchNormalization

Bases: BatchNormalization

Fixed default name of BatchNormalization to match TpuBatchNormalization.

Source code in deeplabcut/pose_estimation_tensorflow/nnets/utils.py

class BatchNormalization(tf.compat.v1.layers.BatchNormalization):
    """Fixed default name of BatchNormalization to match TpuBatchNormalization."""

    def __init__(self, name="tpu_batch_normalization", **kwargs):
        super().__init__(name=name, **kwargs)

DepthwiseConv2D

Bases: DepthwiseConv2D, Layer

Wrap keras DepthwiseConv2D to tf.layers.

Source code in deeplabcut/pose_estimation_tensorflow/nnets/utils.py

class DepthwiseConv2D(tf.keras.layers.DepthwiseConv2D, tf.compat.v1.layers.Layer):
    """Wrap keras DepthwiseConv2D to tf.layers."""

    pass

TpuBatchNormalization

Bases: BatchNormalization

Cross replica batch normalization.

Source code in deeplabcut/pose_estimation_tensorflow/nnets/utils.py

class TpuBatchNormalization(tf.compat.v1.layers.BatchNormalization):
    """Cross replica batch normalization."""

    def __init__(self, fused=False, **kwargs):
        if fused in (True, None):
            raise ValueError("TpuBatchNormalization does not support fused=True.")
        super().__init__(fused=fused, **kwargs)

    @staticmethod
    def _cross_replica_average(t, num_shards_per_group):
        """Calculates the average value of input tensor across TPU replicas."""
        num_shards = tpu_function.get_tpu_context().number_of_shards
        group_assignment = None
        if num_shards_per_group > 1:
            if num_shards % num_shards_per_group != 0:
                raise ValueError(f"num_shards: {num_shards} mod shards_per_group: {num_shards_per_group}, should be 0")
            num_groups = num_shards // num_shards_per_group
            group_assignment = [
                [x for x in range(num_shards) if x // num_shards_per_group == y] for y in range(num_groups)
            ]
        return tpu_ops.cross_replica_sum(t, group_assignment) / tf.cast(num_shards_per_group, t.dtype)

    def _moments(self, inputs, reduction_axes, keep_dims):
        """Compute the mean and variance: it overrides the original _moments."""
        shard_mean, shard_variance = super()._moments(inputs, reduction_axes, keep_dims=keep_dims)

        num_shards = tpu_function.get_tpu_context().number_of_shards or 1
        if num_shards <= 8:  # Skip cross_replica for 2x2 or smaller slices.
            num_shards_per_group = 1
        else:
            num_shards_per_group = max(8, num_shards // 8)
        tf.compat.v1.logging.info(f"TpuBatchNormalization with num_shards_per_group {num_shards_per_group}")
        if num_shards_per_group > 1:
            # Compute variance using: Var[X]= E[X^2] - E[X]^2.
            shard_square_of_mean = tf.math.square(shard_mean)
            shard_mean_of_square = shard_variance + shard_square_of_mean
            group_mean = self._cross_replica_average(shard_mean, num_shards_per_group)
            group_mean_of_square = self._cross_replica_average(shard_mean_of_square, num_shards_per_group)
            group_variance = group_mean_of_square - tf.math.square(group_mean)
            return group_mean, group_variance
        return shard_mean, shard_variance

build_learning_rate

build_learning_rate(
    initial_lr,
    global_step,
    steps_per_epoch=None,
    lr_decay_type="exponential",
    decay_factor=0.97,
    decay_epochs=2.4,
    total_steps=None,
    warmup_epochs=5,
)

Build learning rate.

Source code in deeplabcut/pose_estimation_tensorflow/nnets/utils.py

def build_learning_rate(
    initial_lr,
    global_step,
    steps_per_epoch=None,
    lr_decay_type="exponential",
    decay_factor=0.97,
    decay_epochs=2.4,
    total_steps=None,
    warmup_epochs=5,
):
    """Build learning rate."""
    if lr_decay_type == "exponential":
        assert steps_per_epoch is not None
        decay_steps = steps_per_epoch * decay_epochs
        lr = tf.compat.v1.train.exponential_decay(initial_lr, global_step, decay_steps, decay_factor, staircase=True)
    elif lr_decay_type == "cosine":
        assert total_steps is not None
        lr = 0.5 * initial_lr * (1 + tf.cos(np.pi * tf.cast(global_step, tf.float32) / total_steps))
    elif lr_decay_type == "constant":
        lr = initial_lr
    else:
        raise AssertionError(f"Unknown lr_decay_type : {lr_decay_type}")

    if warmup_epochs:
        tf.compat.v1.logging.info(f"Learning rate warmup_epochs: {warmup_epochs}")
        warmup_steps = int(warmup_epochs * steps_per_epoch)
        warmup_lr = initial_lr * tf.cast(global_step, tf.float32) / tf.cast(warmup_steps, tf.float32)
        lr = tf.cond(
            pred=global_step < warmup_steps,
            true_fn=lambda: warmup_lr,
            false_fn=lambda: lr,
        )

    return lr

build_optimizer

build_optimizer(learning_rate, optimizer_name='rmsprop', decay=0.9, epsilon=0.001, momentum=0.9)

Build optimizer.

Source code in deeplabcut/pose_estimation_tensorflow/nnets/utils.py

def build_optimizer(learning_rate, optimizer_name="rmsprop", decay=0.9, epsilon=0.001, momentum=0.9):
    """Build optimizer."""
    if optimizer_name == "sgd":
        tf.compat.v1.logging.info("Using SGD optimizer")
        optimizer = tf.compat.v1.train.GradientDescentOptimizer(learning_rate=learning_rate)
    elif optimizer_name == "momentum":
        tf.compat.v1.logging.info("Using Momentum optimizer")
        optimizer = tf.compat.v1.train.MomentumOptimizer(learning_rate=learning_rate, momentum=momentum)
    elif optimizer_name == "rmsprop":
        tf.compat.v1.logging.info("Using RMSProp optimizer")
        optimizer = tf.compat.v1.train.RMSPropOptimizer(learning_rate, decay, momentum, epsilon)
    else:
        tf.compat.v1.logging.fatal(f"Unknown optimizer: {optimizer_name}")
    return optimizer

drop_connect

drop_connect(inputs, is_training, drop_connect_rate)

Apply drop connect.

Source code in deeplabcut/pose_estimation_tensorflow/nnets/utils.py

def drop_connect(inputs, is_training, drop_connect_rate):
    """Apply drop connect."""
    if not is_training:
        return inputs

    # Compute keep_prob
    # TODO(tanmingxing): add support for training progress.
    keep_prob = 1.0 - drop_connect_rate

    # Compute drop_connect tensor
    batch_size = tf.shape(input=inputs)[0]
    random_tensor = keep_prob
    random_tensor += tf.random.uniform([batch_size, 1, 1, 1], dtype=inputs.dtype)
    binary_tensor = tf.floor(random_tensor)
    output = tf.compat.v1.div(inputs, keep_prob) * binary_tensor
    return output