def test_ghostbatchnormalization(mode: bool,
virtual_batch_size: Optional[int]) -> None:
# setup up GhostBatchNormalization layer
ghost_batch_norm = GhostBatchNormalization(
OUTPUT_SIZE, virtual_batch_size=virtual_batch_size)
# set training or eval mode
ghost_batch_norm.train(mode=mode)
# setup inputs to test
inputs = torch.randn([BATCH_SIZE, OUTPUT_SIZE], dtype=torch.float32)
# run tensor through
norm_tensor = ghost_batch_norm(inputs)
# check for correctness of output
assert isinstance(norm_tensor, torch.Tensor)
assert norm_tensor.shape == (BATCH_SIZE, OUTPUT_SIZE)
# check for required properties
assert ghost_batch_norm.input_shape == inputs.shape[1:]
assert ghost_batch_norm.output_shape == inputs.shape[1:]
assert ghost_batch_norm.input_dtype == torch.float32
assert isinstance(ghost_batch_norm.moving_mean, torch.Tensor)
assert ghost_batch_norm.moving_mean.shape == (OUTPUT_SIZE, )
assert isinstance(ghost_batch_norm.moving_variance, torch.Tensor)
assert ghost_batch_norm.moving_variance.shape == (OUTPUT_SIZE, )
def __init__(
self,
input_size: int,
size: int,
apply_glu: bool = True,
bn_momentum: float = 0.9,
bn_epsilon: float = 1e-3,
bn_virtual_bs: int = None,
shared_fc_layer: LudwigModule = None,
):
super().__init__()
self.input_size = input_size
self.apply_glu = apply_glu
self.size = size
units = size * 2 if apply_glu else size
if shared_fc_layer:
self.fc_layer = shared_fc_layer
else:
self.fc_layer = torch.nn.Linear(input_size, units, bias=False)
self.batch_norm = GhostBatchNormalization(
units,
virtual_batch_size=bn_virtual_bs,
momentum=bn_momentum,
epsilon=bn_epsilon)
def __init__(
self,
input_size: int,
size: int,
apply_glu: bool = True,
bn_momentum: float = 0.1,
bn_epsilon: float = 1e-3,
bn_virtual_bs: int = None,
shared_fc_layer: LudwigModule = None,
):
super().__init__()
self.input_size = input_size
self.apply_glu = apply_glu
self.size = size
units = size * 2 if apply_glu else size
# Initialize fc_layer before assigning to shared layer for torchscript compatibilty
self.fc_layer = nn.Linear(input_size, units, bias=False)
if shared_fc_layer is not None:
assert shared_fc_layer.weight.shape == self.fc_layer.weight.shape
self.fc_layer = shared_fc_layer
self.batch_norm = GhostBatchNormalization(
units,
virtual_batch_size=bn_virtual_bs,
momentum=bn_momentum,
epsilon=bn_epsilon)
def __init__(
self,
size: int,
apply_glu: bool = True,
bn_momentum: float = 0.9,
bn_virtual_divider: int = 32,
shared_fc_layer: tf.keras.layers.Layer = None,
epsilon: float = 1e-5,
):
super(FeatureBlock, self).__init__()
self.apply_glu = apply_glu
self.size = size
units = size * 2 if apply_glu else size
if shared_fc_layer:
self.fc_layer = shared_fc_layer
else:
self.fc_layer = tf.keras.layers.Dense(units, use_bias=False)
self.batch_norm = GhostBatchNormalization(
virtual_divider=bn_virtual_divider, momentum=bn_momentum
)
def __init__(
self,
size: int,
apply_glu: bool = True,
bn_momentum: float = 0.9,
bn_epsilon: float = 1e-3,
bn_virtual_bs: int = None,
shared_fc_layer: tf.keras.layers.Layer = None,
):
super().__init__()
self.apply_glu = apply_glu
self.size = size
units = size * 2 if apply_glu else size
if shared_fc_layer:
self.fc_layer = shared_fc_layer
else:
self.fc_layer = tf.keras.layers.Dense(units, use_bias=False)
self.batch_norm = GhostBatchNormalization(
virtual_batch_size=bn_virtual_bs,
momentum=bn_momentum,
epsilon=bn_epsilon)