test_case.py

import numpy as np
import tensorflow as tf
import sugartensor as stf

# set log level to debug
stf.sg_verbosity(10)


@stf.sg_layer_func
def seq_dense(tensor, opt):
    w = stf.sg_initializer.he_uniform('W', (1, opt.in_dim, opt.dim))
    b = stf.sg_initializer.constant('b', opt.dim) if opt.bias else 0

    w = tf.squeeze(w, axis=0)

    out = tf.matmul(tensor, w) + b

    return out


#
# setup dataset
#
dataset = [
    ('four five two#____', '452#______________'),
    ('one nine#_________', '19#_______________'),
    ('four eight four#__', '484#______________'),
    ('two four five#____', '245#______________'),
    ('seven one#________', '71#_______________'),
    ('seven five eight#_', '758#______________'),
    ('zero nine nine#___', '099#______________'),
    ('six nine#_________', '69#_______________'),
    ('three nine#_______', '39#_______________'),
    ('four one#_________', '41#_______________')
]
vocabulary = '_#zerontwhufivsxg0123456789 '
encoding = {
    char: code for code, char in enumerate(vocabulary)
}
dataset = [
    ([encoding[char] for char in source], [encoding[char] for char in target])
    for (source, target) in dataset
]
source, target = zip(*dataset)

#
# setup batch queue
#
source = tf.convert_to_tensor(np.asarray(source))
target = tf.convert_to_tensor(np.asarray(target))

source, target = tf.train.slice_input_producer([source, target], shuffle=True)

source, target = tf.train.shuffle_batch(
    [source, target], 10,
    capacity=640,
    min_after_dequeue=320
)

#
# setup loss
#
loss = 0

with tf.name_scope(None, "semi-supervied-x", values=[source, target]):
    x = tf.cast(source, tf.int32)
    y = tf.cast(target, tf.int32)

    emb_x = stf.sg_emb(
        name='embedding-source',
        voca_size=len(vocabulary),
        dim=20
    )
    enc = x.sg_lookup(emb=emb_x)

    with tf.variable_scope("translator", values=[x]):
        logits = enc.sg_conv1d(
            size=1, dim=len(vocabulary),
            name='logits-dense'
        )

    with tf.variable_scope("translator", values=[x], reuse=True):
        # BUG: seq_dense is not passed reuse directly, instead the parent
        # variable_scope has reuse=True
        logits = tf.scan(
            lambda acc, enc_t: seq_dense(enc_t, dim=len(vocabulary),
                                         name='logits-dense'),
            elems=tf.transpose(enc, perm=[1, 0, 2]),
            initializer=tf.zeros(
                (tf.shape(enc)[0], len(vocabulary)), dtype=enc.dtype
            )
        )
        seq_logits = tf.transpose(logits, perm=[1, 0, 2])

    loss += tf.reduce_mean(logits.sg_ce(target=y, mask=True), axis=0)
    loss += tf.reduce_mean(seq_logits.sg_ce(target=y, mask=True), axis=0)

#
# train
#
stf.sg_train(log_interval=30,
             loss=loss,
             ep_size=1,
             max_ep=10000,
             early_stop=False)