def test_DiagGaussianPd():
reset_session()
# Construct the layer
from rinokeras.core.v1x.common.distributions import DiagGaussianPd
layer = DiagGaussianPd(action_shape=(4, 4))
# Make sure that the layer is not None
assert layer is not None
# Encoded values
input_tensor, _ = random_tensor((16, 4, 4))
actions_tensor, _ = random_tensor((16, 4, 4))
# Get the output of the layer
value = layer(input_tensor)
logp_actions = layer.logp_actions(actions_tensor)
entropy = layer.entropy()
# Create a named temporary file for save/restore testing
with tempfile.TemporaryFile() as weights_file:
# Construct the session
output = run_simple_session_save_weights(
inputs=[value, logp_actions, entropy],
feed={},
weights=[layer],
weights_file=weights_file)
# Make sure the value is not null
assert output[0] is not None
assert output[1] is not None
assert output[2] is not None
# Make sure the output shape is correct
assert output[0].shape == (16, 4, 4)
assert output[1].shape == (16, )
assert output[2].shape == ()
# Make sure the output values are correct (If Possible)
pass
# Check loading and restoring
load_restore_test(output=output,
inputs=[value, logp_actions, entropy],
feed={},
weights=[layer],
weights_file=weights_file)
# Do regression testing
check_regression('DiagGaussianPd_expected_output',
output,
__file__,
'regression_outputs/test_distribution_outputs.json',
debug=_RK_REBUILD_REGRESSION)
def test_transformer_decoder():
reset_session()
# Construct the layer
from rinokeras.core.v1x.models.transformer import TransformerDecoder, TransformerInputEmbedding
tie = TransformerInputEmbedding(128, False)
layer = TransformerDecoder(embedding_layer=tie,
output_layer=tf.keras.layers.Dense(128),
n_layers=2,
n_heads=4,
d_model=128,
d_filter=32)
# Make sure that the layer is not None
assert layer is not None
# Encoded values
source_tensor, _ = random_tensor((2,32,128))
target_tensor, _ = random_tensor((2,32,64))
# Get random masking values
source_mask, _ = random_mask_tensor(2, 32)
target_mask, _ = random_mask_tensor(2, 32)
source_mask = convert_to_attention_mask(source_tensor, source_mask)
target_mask = convert_to_attention_mask(target_tensor, target_mask)
# Get the output of the layer
value = layer((source_tensor, target_tensor), mask=(source_mask, target_mask))
# Create a named temporary file for save/restore testing
with tempfile.TemporaryFile() as weights_file:
# Construct the session
output = run_simple_session_save_weights(inputs=[value],
feed={},
weights=[layer],
weights_file=weights_file)
assert_not_none(output)
assert_expected_shapes(output, [(2,32,128)])
# Check loading and restoring
load_restore_test(output=output,
inputs=[value],
feed={},
weights=[layer],
weights_file=weights_file)
# Do regression testing
check_regression('transformer_decoder_output',
output, __file__, 'regression_outputs/test_transformer_decoder_outputs.json', debug=_RK_REBUILD_REGRESSION)
# Do a config test
from_config_test(TransformerDecoder, layer)
def test_transformer_multi_attention():
reset_session()
# Construct the layer
from rinokeras.core.v1x.models.transformer.transformer_attention import TransformerMultiAttention
layer = TransformerMultiAttention(n_heads=4)
# Make sure that the layer is not None
assert layer is not None
# Encoded values
encoder_tensor, _ = random_tensor((16, 32, 64))
decoder_tensor, _ = random_tensor((16, 32, 64))
# Get the output of the layer
value = layer((encoder_tensor, decoder_tensor))
# Create a named temporary file for save/restore testing
with tempfile.TemporaryFile() as weights_file:
# Construct the session
output = run_simple_session_save_weights(inputs=[value],
feed={},
weights=[layer],
weights_file=weights_file)
assert_not_none(output)
assert_expected_shapes(output, [(16, 32, 64)])
# Check loading and restoring
load_restore_test(output=output,
inputs=[value],
feed={},
weights=[layer],
weights_file=weights_file)
# Do regression testing
check_regression(
'transformer_multi_attention_expected_output',
output,
__file__,
'regression_outputs/test_transformer_attention_outputs.json',
debug=_RK_REBUILD_REGRESSION)
# Do a config test
from_config_test(TransformerMultiAttention, layer)
def test_transformer_decoder_block_no_mask():
reset_session()
# Construct the layer
from rinokeras.core.v1x.models.transformer import TransformerDecoderBlock
layer = TransformerDecoderBlock(n_heads=4, filter_size=128, hidden_size=64)
# Make sure that the layer is not None
assert layer is not None
# Encoded values
source_tensor, _ = random_tensor((2,32,64))
target_tensor, _ = random_tensor((2,32,64))
# Get the output of the layer
_, value = layer((source_tensor, target_tensor))
# Create a named temporary file for save/restore testing
with tempfile.TemporaryFile() as weights_file:
# Construct the session
output = run_simple_session_save_weights(inputs=[value],
feed={},
weights=[layer],
weights_file=weights_file)
assert_not_none(output)
assert_expected_shapes(output, [(2,32,64)])
# Check loading and restoring
load_restore_test(output=output,
inputs=[value],
feed={},
weights=[layer],
weights_file=weights_file)
# Do regression testing
check_regression('transformer_decoder_block_no_mask_output',
output, __file__, 'regression_outputs/test_transformer_decoder_outputs.json', debug=_RK_REBUILD_REGRESSION)
# Do a config test
from_config_test(TransformerDecoderBlock, layer)
def test_transformer_encoder_masking_with_conv():
reset_session()
# Construct the layer
from rinokeras.core.v1x.models.transformer import TransformerEncoder
layer = TransformerEncoder(embedding_layer=tf.keras.layers.Dense(64),
n_layers=2,
n_heads=4,
d_model=64,
d_filter=128)
# Make sure that the layer is not None
assert layer is not None
# Encoded values
input_tensor, _ = random_tensor((16, 32, 128))
input_mask, _ = random_mask_tensor(16, 32)
input_mask = convert_to_attention_mask(input_tensor, input_mask)
conv_mask, _ = random_mask_tensor(16, 32)
conv_mask = convert_sequence_length_to_sequence_mask(
input_tensor, conv_mask)
# Get the output of the layer
value = layer(input_tensor, mask=(input_mask, conv_mask))
# Create a named temporary file for save/restore testing
with tempfile.TemporaryFile() as weights_file:
# Construct the session
output = run_simple_session_save_weights(inputs=[value],
feed={},
weights=[layer],
weights_file=weights_file)
assert_not_none(output)
assert_expected_shapes(output, [(16, 32, 64)])
# Check loading and restoring
load_restore_test(output=output,
inputs=[value],
feed={},
weights=[layer],
weights_file=weights_file)
# Do regression testing
check_regression(
'transformer_encoder_output_masking_with_conv',
output,
__file__,
'regression_outputs/test_transformer_encoder_outputs.json',
debug=_RK_REBUILD_REGRESSION)
# Do a config test
from_config_test(TransformerEncoder, layer)
def test_transformer_decoder_fast_beam_decode_discrete():
reset_session()
# Construct the layer
from rinokeras.core.v1x.models.transformer import TransformerDecoder, TransformerInputEmbedding
tie = TransformerInputEmbedding(128, discrete=True, n_symbols=256)
layer = TransformerDecoder(embedding_layer=tie,
output_layer=tf.keras.layers.Dense(256),
n_layers=2,
n_heads=4,
d_model=128,
d_filter=32)
# Make sure that the layer is not None
assert layer is not None
# Encoded values
source_tensor, _ = random_tensor((2,32,128))
# Get the output of the layer
value, scores = layer.fast_beam_decode(source_tensor, 20, batch_size=2, n_beams=4)
# Create a named temporary file for save/restore testing
with tempfile.TemporaryFile() as weights_file:
# Construct the session
output = run_simple_session_save_weights(inputs=[value, scores],
feed={},
weights=[layer],
weights_file=weights_file)
assert_not_none(output)
assert_expected_shapes(output, [(2,4, 20), (2,4)])
# Check loading and restoring
load_restore_test(output=output,
inputs=[value],
feed={},
weights=[layer],
weights_file=weights_file)
# Do regression testing
check_regression('transformer_decoder_fast_beam_decode',
output, __file__, 'regression_outputs/test_transformer_decoder_outputs.json', debug=_RK_REBUILD_REGRESSION)
# Do a config test
from_config_test(TransformerDecoder, layer)
def test_transformer_input_embedding_non_discrete():
reset_session()
# Construct the layer
from rinokeras.core.v1x.models.transformer import TransformerInputEmbedding
layer = TransformerInputEmbedding(embed_size=128, discrete=False)
# Make sure that the layer is not None
assert layer is not None
# Encoded values
input_tensor, _ = random_tensor((16, 32, 64))
# Get the output of the layerembedding_layer
value = layer(input_tensor)
# Create a named temporary file for save/restore testing
with tempfile.TemporaryFile() as weights_file:
# Construct the session
output = run_simple_session_save_weights(inputs=[value],
feed={},
weights=[layer],
weights_file=weights_file)
assert_not_none(output)
assert_expected_shapes(output, [(16, 32, 128)])
# Check loading and restoring
load_restore_test(output=output,
inputs=[value],
feed={},
weights=[layer],
weights_file=weights_file)
# Do regression testing
check_regression(
'transformer_input_embedding_non_discrete_expected_output',
output,
__file__,
'regression_outputs/test_transformer_embedding_outputs.json',
debug=_RK_REBUILD_REGRESSION)
# Do a config test
from_config_test(TransformerInputEmbedding, layer)
def test_normed_conv_stack_3d():
reset_session()
# Construct the layer
from rinokeras.core.v1x.common.layers.conv import NormedConvStack
ncs_layer = NormedConvStack(dimension=3, filters=12, kernel_size=4)
# Make sure that the layer is not None
assert ncs_layer is not None
# Encoded values
input_tensor, input_values = random_tensor((2, 8, 16, 16, 16))
# Get the output of the layer
value = ncs_layer(input_tensor)
# Create a named temporary file for save/restore testing
with tempfile.TemporaryFile() as weights_file:
# Construct the session
output = run_simple_session_save_weights(inputs=[value],
feed={},
weights=[ncs_layer],
weights_file=weights_file)
# Make sure the value is not null
assert output[0] is not None
# Make sure the output shape is correct
assert output[0].shape == (2, 8, 16, 16, 12)
# Make sure the output values are correct (If Possible)
pass
# Check loading and restoring
load_restore_test(output=output,
inputs=[value],
feed={},
weights=[ncs_layer],
weights_file=weights_file)
# Do regression testing
check_regression('normed_conv_stack_3d_expected_output',
output, __file__, 'regression_outputs/test_conv_outputs.json', debug=_RK_REBUILD_REGRESSION)
def test_qanet_conv_block():
reset_session()
# Construct the layer
from rinokeras.core.v1x.models.qanet import QANetConvBlock
layer = QANetConvBlock(filters=64, kernel_size=7)
# Make sure that the layer is not None
assert layer is not None
# Encoded values
input_tensor, _ = random_tensor((16,32,64))
# Get the output of the layer
value = layer(input_tensor)
# Create a named temporary file for save/restore testing
with tempfile.TemporaryFile() as weights_file:
# Construct the session
output = run_simple_session_save_weights(inputs=[value],
feed={},
weights=[layer],
weights_file=weights_file)
assert_not_none(output)
assert_expected_shapes(output, [(16,32,64)])
# Check loading and restoring
load_restore_test(output=output,
inputs=[value],
feed={},
weights=[layer],
weights_file=weights_file)
# Do regression testing
check_regression('qanet_conv_block_expected_output',
output, __file__, 'regression_outputs/test_qanet_ff_outputs.json', debug=_RK_REBUILD_REGRESSION)
# Do a config test
from_config_test(QANetConvBlock, layer)
def test_grouped_conv():
reset_session()
# Construct the layer
from rinokeras.core.v1x.common.layers.conv import GroupedConvolution
rb_layer = GroupedConvolution()
# Make sure that the layer is not None
assert rb_layer is not None
# Encoded values
input_tensor, _ = random_tensor((2, 16, 16, 3))
# Get the output of the layer
value = rb_layer(input_tensor)
# Create a named temporary file for save/restore testing
with tempfile.TemporaryFile() as weights_file:
# Construct the session
output = run_simple_session_save_weights(inputs=[value],
feed={},
weights=[rb_layer],
weights_file=weights_file)
# Make sure the value is not null
assert output[0] is not None
# Make sure the output shape is correct
assert output[0].shape == (2, 16, 16, 64)
# Check loading and restoring
load_restore_test(output=output,
inputs=[value],
feed={},
weights=[rb_layer],
weights_file=weights_file)
# Do regression testing
check_regression('grouped_conv_expected_output',
output, __file__, 'regression_outputs/test_conv_outputs.json', debug=_RK_REBUILD_REGRESSION)
def test_resnext50_base():
reset_session()
# Construct the layer
from rinokeras.core.v1x.models.resnet.resnet import ResNeXt50
layer = ResNeXt50()
# Make sure that the layer is not None
assert layer is not None
# Encoded values
inputs_tensor, _ = random_tensor((2, 8, 8, 3))
# Get the output of the layer
value = layer(inputs_tensor)
# Create a named temporary file for save/restore testing
with tempfile.TemporaryFile() as weights_file:
# Construct the session
output = run_simple_session_save_weights(inputs=[value],
feed={},
weights=[layer],
weights_file=weights_file)
assert_not_none(output)
assert_expected_shapes(output, [(2, 2048)])
# Check loading and restoring
load_restore_test(output=output,
inputs=[value],
feed={},
weights=[layer],
weights_file=weights_file)
# Do regression testing
check_regression('resnext50_base_expected_output',
output,
__file__,
'regression_outputs/test_resnet_outputs.json',
debug=_RK_REBUILD_REGRESSION)
