import numpy as np
import tensorflow as tf
from sleap.nn.system import use_cpu_only; use_cpu_only() # hide GPUs for test
from sleap.nn.architectures import leap
from sleap.nn.config import LEAPConfig
class LeapTests(tf.test.TestCase):
def test_leap_cnn_reference(self):
# Reference implementation from the original paper.
arch = leap.LeapCNN(
filters=64,
filters_rate=2,
down_blocks=3,
down_convs_per_block=3,
up_blocks=3,
up_interpolate=False,
up_convs_per_block=2,
)
x_in = tf.keras.layers.Input((192, 192, 1))
x, x_mid = arch.make_backbone(x_in)
model = tf.keras.Model(x_in, x)
param_counts = [
np.prod(train_var.shape) for train_var in model.trainable_weights
]
with self.subTest("number of layers"):
self.assertEqual(len(model.layers), 40)
with self.subTest("number of trainable weights"):
self.assertEqual(len(model.trainable_weights), 36)
with self.subTest("trainable parameter count"):
import numpy as np
import tensorflow as tf
from sleap.nn.system import use_cpu_only
use_cpu_only() # hide GPUs for test
from sleap.nn.architectures import encoder_decoder
class EncoderDecoderTests(tf.test.TestCase):
def test_simple_conv_block(self):
block = encoder_decoder.SimpleConvBlock(
pooling_stride=2,
num_convs=3,
filters=16,
kernel_size=3,
use_bias=True,
batch_norm=False,
batch_norm_before_activation=True,
activation="relu",
)
x_in = tf.keras.Input((8, 8, 1))
x = block.make_block(x_in)
model = tf.keras.Model(x_in, x)
self.assertEqual(len(model.layers), 1 + 2 * 3 + 1)
self.assertEqual(len(model.trainable_weights), 6)
self.assertEqual(model.count_params(), 4800)
self.assertAllEqual(model.output.shape, (None, 4, 4, 16))
def test_simple_conv_block_bn(self):
