def test_encoder_embedding_build(self, batch_size, feature_dim, scope):
os.chdir(os.path.dirname(os.path.realpath(__file__)))
os.chdir('../')
from model import AMConfig
from model import ConformerAM
config_file = 'configs/train_fp32_kl_loss.yaml'
config = AMConfig.from_yaml(config_file)
config['use_ipu_dropout'] = False
model = ConformerAM(config)
from util import get_config
cfg = get_config(1)
cfg.configure_ipu_system()
input_len = tf.constant([
model.config['maxlen_in'],
])
input = tf.random_normal(shape=(batch_size, model.config['maxlen_in'],
feature_dim, 1))
encoder_embedding = model._build_encoder_embedding(
input, input_len, scope)
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
x, pos_emb, mask_adder = sess.run(encoder_embedding)
assert x.shape == (batch_size, model.config['maxlen_in'] // 4 - 1,
model.config['adim'])
assert pos_emb.shape == (1, (model.config['maxlen_in'] // 4 - 1),
model.config['adim'])
assert mask_adder.shape == (1, 1, 1,
model.config['maxlen_in'] // 4 - 1)
def test_decoder_layer_build(self, batch_size, prefix):
os.chdir(os.path.dirname(os.path.realpath(__file__)))
os.chdir('../')
from model import AMConfig
from model import ConformerAM
from util import get_config
cfg = get_config(1)
cfg.configure_ipu_system()
config_file = 'configs/train_fp32_kl_loss.yaml'
config = AMConfig.from_yaml(config_file)
config['use_ipu_dropout'] = False
model = ConformerAM(config)
tgt = tf.ones(shape=(batch_size, model.config['maxlen_tgt'],
model.config['adim']))
tgt_mask = tf.random_normal(shape=(batch_size, 1,
model.config['maxlen_tgt'],
model.config['maxlen_tgt']))
mem = tf.random_normal(shape=(batch_size, model.config['maxlen_in'],
model.config['adim']))
mem_mask = tf.random_normal(shape=(batch_size, 1, 1,
model.config['maxlen_in']))
layer = model._build_decoder_layer(tgt, tgt_mask, mem, mem_mask,
prefix)
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
output = sess.run(layer)
assert output.shape == (batch_size, model.config['maxlen_tgt'],
model.config['adim'])
def test_load_yaml_config(self):
os.chdir(os.path.dirname(os.path.realpath(__file__)))
os.chdir('../')
from model import AMConfig
config_file = 'configs/train_fp32_kl_loss.yaml'
config = AMConfig.from_yaml(config_file)
assert config['batch_size']
def test_build_computation_stages(self):
os.chdir(os.path.dirname(os.path.realpath(__file__)))
os.chdir('../')
from model import AMConfig
from model import ConformerAM
from util import get_config
cfg = get_config(1)
cfg.configure_ipu_system()
config_file = 'configs/train_fp32_kl_loss.yaml'
config = AMConfig.from_yaml(config_file)
model = ConformerAM(config)
model._build_computational_stages()
def test_pos_embedding_build(self, maxlen, dmodel):
os.chdir(os.path.dirname(os.path.realpath(__file__)))
os.chdir('../')
from model import AMConfig
from model import ConformerAM
from util import get_config
cfg = get_config(1)
cfg.configure_ipu_system()
config_file = 'configs/train_fp32_kl_loss.yaml'
config = AMConfig.from_yaml(config_file)
model = ConformerAM(config)
pos_emb = model._build_pos_embedding(maxlen, dmodel)
batch_size = 1
with tf.Session() as sess:
res = sess.run(pos_emb)
assert res.shape == (batch_size, maxlen, dmodel)
def test_get_pipeline_depth(self):
os.chdir(os.path.dirname(os.path.realpath(__file__)))
os.chdir('../')
from model import AMConfig
from model import ConformerAM
from util import get_config
cfg = get_config(1)
cfg.configure_ipu_system()
config_file = 'configs/train_fp32_kl_loss.yaml'
config = AMConfig.from_yaml(
config_file, **{
'data_path': './train',
'dict_path': './sample_train_units.txt',
'use_synthetic_data': True
})
model = ConformerAM(config)
model._build_dataset()
model._build_computational_stages()
def test_relative_shift_build(self, batch_size, n_head, length_q,
length_v):
os.chdir(os.path.dirname(os.path.realpath(__file__)))
os.chdir('../')
from model import AMConfig
from model import ConformerAM
from util import get_config
cfg = get_config(1)
cfg.configure_ipu_system()
config_file = 'configs/train_fp32_kl_loss.yaml'
config = AMConfig.from_yaml(config_file)
config['use_ipu_dropout'] = False
model = ConformerAM(config)
input = tf.ones(shape=(batch_size, n_head, length_q, length_v))
shift = model._relative_shift(input)
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
output = sess.run(shift)
assert output.shape == (batch_size, n_head, length_q, length_v)
def test_conv_build(self, batch_size, length, scope_name):
os.chdir(os.path.dirname(os.path.realpath(__file__)))
os.chdir('../')
from model import AMConfig
from model import ConformerAM
from util import get_config
cfg = get_config(1)
cfg.configure_ipu_system()
config_file = 'configs/train_fp32_kl_loss.yaml'
config = AMConfig.from_yaml(config_file)
config['use_ipu_dropout'] = False
model = ConformerAM(config)
input = tf.random_normal(shape=(batch_size, length,
model.config['adim']))
conv = model._build_conv_module(input, scope_name)
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
output = sess.run(conv)
assert output.shape == (batch_size, length, model.config['adim'])
def test_self_attention_build(self, batch_size, lq, lv, scope_name):
os.chdir(os.path.dirname(os.path.realpath(__file__)))
os.chdir('../')
from model import AMConfig
from model import ConformerAM
from util import get_config
cfg = get_config(1)
cfg.configure_ipu_system()
config_file = 'configs/train_fp32_kl_loss.yaml'
config = AMConfig.from_yaml(config_file)
config['use_ipu_dropout'] = False
model = ConformerAM(config)
query = tf.random_normal(shape=(batch_size, lq, model.config['adim']))
key = tf.random_normal(shape=(batch_size, lv, model.config['adim']))
value = tf.random_normal(shape=(batch_size, lv, model.config['adim']))
self_attention = model._build_self_attention(query, key, value,
scope_name)
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
output = sess.run(self_attention)
assert output.shape == (batch_size, lq, model.config['adim'])
type=str,
default='model.log',
help="log file for model parameters")
parser.add_argument('--seed',
type=int,
default=1991,
help='set random seed for all')
parser.add_argument('--data-path',
type=str,
default='data/train',
help='set training data path')
parser.add_argument('--dict-path',
type=str,
default='./sample_train_units.txt',
help='set training data vocab path')
parser.add_argument('--use-synthetic-data',
type=str,
default=True,
help='if use synthetic data')
parser.add_argument(
'--wandb-name',
type=str,
default=None,
help='Name of the Weights&Biases run, disabled if None')
args = parser.parse_args()
seed_all(args.seed)
config = AMConfig.from_yaml(args.config, **vars(args))
conformeram = ConformerAM(config)
conformeram.run_with_pipeline()