Exemplo n.º 1
0
def train_mul(Model):

    # create dataset and dataloader
    with tf.device("/cpu:0"):
        tfdata_dev = TFData(dataset=None,
                            dataAttr=['feature', 'label', 'align'],
                            dir_save=args.dirs.dev.tfdata,
                            args=args).read(_shuffle=False)
        tfdata_monitor = TFData(dataset=None,
                                dataAttr=['feature', 'label', 'align'],
                                dir_save=args.dirs.train.tfdata,
                                args=args).read(_shuffle=False)
        tfdata_monitor = tfdata_monitor.repeat().shuffle(500).padded_batch(
            args.batch_size,
            ([None, args.dim_input], [None], [None])).prefetch(buffer_size=5)
        tfdata_iter = iter(tfdata_monitor)
        tfdata_dev = tfdata_dev.padded_batch(
            args.batch_size, ([None, args.dim_input], [None], [None]))

    # get dataset ngram
    ngram_py, total_num = read_ngram(args.data.k,
                                     args.dirs.ngram,
                                     args.token2idx,
                                     type='list')

    # create model paremeters
    opti = tf.keras.optimizers.SGD(0.5)
    model = Model(args, optimizer=opti, name='fc')
    model.summary()

    def thread_session(thread_id, queue_input, queue_output):
        global kernel
        gpu = args.list_gpus[thread_id]
        with tf.device(gpu):
            opti_adam = build_optimizer(args, type='adam')
            model = Model(args,
                          optimizer=opti_adam,
                          name='fc' + str(thread_id))
            print('thread_{} is waiting to run on {}....'.format(
                thread_id, gpu))
            while True:
                # s = time()
                id, weights, x, aligns_sampled = queue_input.get()
                model.set_weights(weights)
                # t = time()
                logits = model(x, training=False)
                pz, K = model.EODM(logits, aligns_sampled, kernel)
                queue_output.put((id, pz, K))
                # print('{} {:.3f}|{:.3f}s'.format(gpu, t-s, time()-s))

    for id in range(args.num_gpus):
        thread = threading.Thread(target=thread_session,
                                  args=(id, queue_input, queue_output))
        thread.daemon = True
        thread.start()

    # save & reload
    ckpt = tf.train.Checkpoint(model=model, optimizer=opti)
    ckpt_manager = tf.train.CheckpointManager(ckpt,
                                              args.dir_checkpoint,
                                              max_to_keep=20)
    if args.dirs.restore:
        latest_checkpoint = tf.train.CheckpointManager(
            ckpt, args.dirs.restore, max_to_keep=1).latest_checkpoint
        ckpt.restore(latest_checkpoint)
        print('{} restored!!'.format(latest_checkpoint))

    best_rewards = -999
    start_time = datetime.now()
    fer = 1.0
    seed = 99999
    step = 0
    global aligns_sampled, kernel

    while 1:
        if fer < 0.69:
            break
        elif fer > 0.77 or step > 69:
            print('{}-th reset, pre FER: {:.3f}'.format(seed, fer))
            seed += 1
            step = 0
            tf.random.set_seed(seed)
            model = Model(args, optimizer=opti, name='fc')
            head_tail_constrain(next(tfdata_iter), model, opti)
            fer = mini_eva(tfdata_dev, model)
            ngram_sampled = sample(ngram_py, args.data.top_k)
            kernel, py = ngram2kernel(ngram_sampled, args)
        else:
            step += 1
            loss = train_step(model, tfdata_iter, py)
            fer = mini_eva(tfdata_dev, model)
            print('\tloss: {:.3f}\tFER: {:.3f}'.format(loss, fer))

    for global_step in range(99999):
        run_model_time = time()
        loss = train_step(model, tfdata_iter, py)

        used_time = time() - run_model_time
        if global_step % 1 == 0:
            print('full training loss: {:.3f}, spend {:.2f}s step {}'.format(
                loss, used_time, global_step))

        if global_step % args.dev_step == 0:
            evaluation(tfdata_dev, model)
        if global_step % args.decode_step == 0:
            decode(model)
        if global_step % args.fs_step == 0:
            fs_constrain(next(tfdata_iter), model, opti)
        if global_step % args.save_step == 0:
            ckpt_manager.save()

    print('training duration: {:.2f}h'.format(
        (datetime.now() - start_time).total_seconds() / 3600))
Exemplo n.º 2
0
def Train():
    with tf.device("/cpu:0"):
        # dataset_train = ASR_align_ArkDataSet(
        #     scp_file=args.dirs.train.scp,
        #     trans_file=args.dirs.train.trans,
        #     align_file=None,
        #     feat_len_file=None,
        #     args=args,
        #     _shuffle=False,
        #     transform=False)
        # dataset_dev = ASR_align_ArkDataSet(
        #     scp_file=args.dirs.dev.scp,
        #     trans_file=args.dirs.dev.trans,
        #     align_file=None,
        #     feat_len_file=None,
        #     args=args,
        #     _shuffle=False,
        #     transform=False)
        dataset_train = ASR_align_DataSet(
            trans_file=args.dirs.train.trans,
            align_file=args.dirs.train.align,
            uttid2wav=args.dirs.train.wav_scp,
            feat_len_file=args.dirs.train.feat_len,
            args=args,
            _shuffle=False,
            transform=True)
        dataset_dev = ASR_align_DataSet(trans_file=args.dirs.dev.trans,
                                        uttid2wav=args.dirs.dev.wav_scp,
                                        align_file=args.dirs.dev.align,
                                        feat_len_file=args.dirs.dev.feat_len,
                                        args=args,
                                        _shuffle=False,
                                        transform=True)
        # wav data
        feature_train = TFData(dataset=dataset_train,
                               dir_save=args.dirs.train.tfdata,
                               args=args).read(transform=True)
        feature_dev = TFData(dataset=dataset_dev,
                             dir_save=args.dirs.dev.tfdata,
                             args=args).read(transform=True)
        bucket = tf.data.experimental.bucket_by_sequence_length(
            element_length_func=lambda uttid, x: tf.shape(x)[0],
            bucket_boundaries=args.list_bucket_boundaries,
            bucket_batch_sizes=args.list_batch_size,
            padded_shapes=((), [None, args.dim_input]))
        iter_feature_train = iter(
            feature_train.repeat().shuffle(10).apply(bucket).prefetch(
                buffer_size=5))
        # iter_feature_train = iter(feature_train.repeat().shuffle(500).padded_batch(args.batch_size,
        #         ((), [None, args.dim_input])).prefetch(buffer_size=5))
        # feature_dev = feature_dev.apply(bucket).prefetch(buffer_size=5)
        feature_dev = feature_dev.padded_batch(args.batch_size,
                                               ((), [None, args.dim_input]))

    # create model paremeters
    model = Model(args)
    model.summary()
    # lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(
    #     args.opti.lr,
    #     decay_steps=args.opti.decay_steps,
    #     decay_rate=0.5,
    #     staircase=True)
    optimizer = tf.keras.optimizers.Adam(args.opti.lr, beta_1=0.5, beta_2=0.9)

    writer = tf.summary.create_file_writer(str(args.dir_log))
    ckpt = tf.train.Checkpoint(model=model, optimizer=optimizer)
    ckpt_manager = tf.train.CheckpointManager(ckpt,
                                              args.dir_checkpoint,
                                              max_to_keep=20)
    step = 0

    # if a checkpoint exists, restore the latest checkpoint.
    if args.dirs.checkpoint:
        _ckpt_manager = tf.train.CheckpointManager(ckpt,
                                                   args.dirs.checkpoint,
                                                   max_to_keep=1)
        ckpt.restore(_ckpt_manager.latest_checkpoint)
        print('checkpoint {} restored!!'.format(
            _ckpt_manager.latest_checkpoint))
        step = int(_ckpt_manager.latest_checkpoint.split('-')[-1])

    start_time = datetime.now()
    num_processed = 0
    while step < 99999999:
        start = time()

        uttids, x = next(iter_feature_train)
        trans = dataset_train.get_attrs('trans', uttids.numpy())
        loss_supervise = train_CTC_supervised(x, trans, model, optimizer)

        num_processed += len(x)
        progress = num_processed / args.data.train_size

        if step % 10 == 0:
            print('loss: {:.3f}\tbatch: {}\tused: {:.3f}\t {:.3f}% step: {}'.
                  format(loss_supervise, x.shape,
                         time() - start, progress * 100, step))
            with writer.as_default():
                tf.summary.scalar("costs/loss_supervise",
                                  loss_supervise,
                                  step=step)
        if step % args.dev_step == 0:
            cer = evaluate(feature_dev, dataset_dev, args.data.dev_size, model)
            with writer.as_default():
                tf.summary.scalar("performance/cer", cer, step=step)
        if step % args.decode_step == 0:
            monitor(dataset_dev[0], model)
        if step % args.save_step == 0:
            save_path = ckpt_manager.save(step)
            print('save model {}'.format(save_path))

        step += 1

    print('training duration: {:.2f}h'.format(
        (datetime.now() - start_time).total_seconds() / 3600))
Exemplo n.º 3
0
def Train():
    with tf.device("/cpu:0"):
        dataset_train = ASR_align_ArkDataSet(scp_file=args.dirs.train.scp,
                                             trans_file=args.dirs.train.trans,
                                             align_file=None,
                                             feat_len_file=None,
                                             args=args,
                                             _shuffle=False,
                                             transform=False)
        dataset_dev = ASR_align_ArkDataSet(scp_file=args.dirs.dev.scp,
                                           trans_file=args.dirs.dev.trans,
                                           align_file=None,
                                           feat_len_file=None,
                                           args=args,
                                           _shuffle=False,
                                           transform=False)
        # wav data
        feature_train = TFData(dataset=dataset_train,
                               dir_save=args.dirs.train.tfdata,
                               args=args).read()
        feature_dev = TFData(dataset=dataset_dev,
                             dir_save=args.dirs.dev.tfdata,
                             args=args).read()
        bucket = tf.data.experimental.bucket_by_sequence_length(
            element_length_func=lambda uttid, x: tf.shape(x)[0],
            bucket_boundaries=args.list_bucket_boundaries,
            bucket_batch_sizes=args.list_batch_size,
            padded_shapes=((), [None, args.dim_input]))
        iter_feature_train = iter(
            feature_train.repeat().shuffle(500).apply(bucket).prefetch(
                buffer_size=5))
        # iter_feature_train = iter(feature_train.repeat().shuffle(500).padded_batch(args.batch_size,
        #         ((), [None, args.dim_input])).prefetch(buffer_size=5))
        feature_dev = feature_dev.padded_batch(args.batch_size,
                                               ((), [None, args.dim_input]))

    stratedy = tf.distribute.MirroredStrategy(
        devices=["device:GPU:0", "device:GPU:1"])
    with stratedy.scope():
        # create model paremeters
        model = conv_lstm(args)
        model.summary()
        optimizer = tf.keras.optimizers.Adam(args.opti.lr,
                                             beta_1=0.5,
                                             beta_2=0.9)

        writer = tf.summary.create_file_writer(str(args.dir_log))
        ckpt = tf.train.Checkpoint(model=model, optimizer=optimizer)
        ckpt_manager = tf.train.CheckpointManager(ckpt,
                                                  args.dir_checkpoint,
                                                  max_to_keep=20)
        step = 0

        # if a checkpoint exists, restore the latest checkpoint.
        if args.dirs.checkpoint:
            _ckpt_manager = tf.train.CheckpointManager(ckpt,
                                                       args.dirs.checkpoint,
                                                       max_to_keep=1)
            ckpt.restore(_ckpt_manager.latest_checkpoint)
            print('checkpoint {} restored!!'.format(
                _ckpt_manager.latest_checkpoint))
            step = int(_ckpt_manager.latest_checkpoint.split('-')[-1])

    start_time = datetime.now()
    num_processed = 0
    while step < 99999999:
        start = time()

    @tf.function(experimental_relax_shapes=True)
    def _train():
        for _ in tf.range(1000):
            stratedy.experimental_run_v2(train_CTC_supervised,
                                         args=(next(iter_feature_train),
                                               dataset_train, model,
                                               optimizer))
        # loss_supervise = tf.reduce_mean(res._values)
        # loss_supervise = train_CTC_supervised(x, trans, model, optimizer)

        # num_processed += len(x)
        # progress = num_processed / args.data.train_size

        # if step % 10 == 0:
        #     print('loss: {:.3f}\tused: {:.3f}\t step: {}'.format(
        #            loss_supervise, time()-start, step))
        #     with writer.as_default():
        #         tf.summary.scalar("costs/loss_supervise", loss_supervise, step=step)
        # if step % args.dev_step == 0:
        #     cer = evaluate(feature_dev, dataset_dev, args.data.dev_size, model)
        #     with writer.as_default():
        #         tf.summary.scalar("performance/cer", cer, step=step)
        # if step % args.decode_step == 0:
        #     monitor(dataset_dev[0], model)
        # if step % args.save_step == 0:
        #     save_path = ckpt_manager.save(step)
        #     print('save model {}'.format(save_path))
        #
        # step += 1

    print('training duration: {:.2f}h'.format(
        (datetime.now() - start_time).total_seconds() / 3600))
Exemplo n.º 4
0
def Train():
    args.data.untrain_size = TFData.read_tfdata_info(
        args.dirs.untrain.tfdata)['size_dataset']
    with tf.device("/cpu:0"):
        dataset_train = ASR_align_ArkDataSet(scp_file=args.dirs.train.scp,
                                             trans_file=args.dirs.train.trans,
                                             align_file=None,
                                             feat_len_file=None,
                                             args=args,
                                             _shuffle=False,
                                             transform=False)
        dataset_untrain = ASR_align_ArkDataSet(scp_file=args.dirs.untrain.scp,
                                               trans_file=None,
                                               align_file=None,
                                               feat_len_file=None,
                                               args=args,
                                               _shuffle=False,
                                               transform=False)
        dataset_dev = ASR_align_ArkDataSet(scp_file=args.dirs.dev.scp,
                                           trans_file=args.dirs.dev.trans,
                                           align_file=None,
                                           feat_len_file=None,
                                           args=args,
                                           _shuffle=False,
                                           transform=False)
        # wav data
        feature_train = TFData(dataset=dataset_train,
                               dir_save=args.dirs.train.tfdata,
                               args=args).read()
        feature_unsupervise = TFData(dataset=dataset_untrain,
                                     dir_save=args.dirs.untrain.tfdata,
                                     args=args).read()
        feature_dev = TFData(dataset=dataset_dev,
                             dir_save=args.dirs.dev.tfdata,
                             args=args).read()
        bucket = tf.data.experimental.bucket_by_sequence_length(
            element_length_func=lambda uttid, x: tf.shape(x)[0],
            bucket_boundaries=args.list_bucket_boundaries,
            bucket_batch_sizes=args.list_batch_size,
            padded_shapes=((), [None, args.dim_input]))
        iter_feature_train = iter(
            feature_train.repeat().shuffle(100).apply(bucket).prefetch(
                buffer_size=5))
        # iter_feature_unsupervise = iter(feature_unsupervise.repeat().shuffle(100).apply(bucket).prefetch(buffer_size=5))
        # iter_feature_train = iter(feature_train.repeat().shuffle(100).padded_batch(args.batch_size,
        #         ((), [None, args.dim_input])).prefetch(buffer_size=5))
        iter_feature_unsupervise = iter(
            feature_unsupervise.repeat().shuffle(100).padded_batch(
                args.batch_size,
                ((), [None, args.dim_input])).prefetch(buffer_size=5))
        # feature_dev = feature_dev.apply(bucket).prefetch(buffer_size=5)
        feature_dev = feature_dev.padded_batch(args.batch_size,
                                               ((), [None, args.dim_input]))

        dataset_text = TextDataSet(list_files=[args.dirs.lm.data],
                                   args=args,
                                   _shuffle=True)
        tfdata_train = tf.data.Dataset.from_generator(dataset_text, (tf.int32),
                                                      (tf.TensorShape([None])))
        iter_text = iter(tfdata_train.cache().repeat().shuffle(1000).map(
            lambda x: x[:args.model.D.max_label_len]).padded_batch(
                args.text_batch_size,
                ([args.model.D.max_label_len])).prefetch(buffer_size=5))

    # create model paremeters
    encoder = Encoder(args)
    decoder = Decoder(args)
    D = CLM(args)
    encoder.summary()
    decoder.summary()
    D.summary()
    optimizer = tf.keras.optimizers.Adam(0.0001, beta_1=0.5, beta_2=0.9)
    optimizer_D = tf.keras.optimizers.Adam(0.0001, beta_1=0.5, beta_2=0.9)

    writer = tf.summary.create_file_writer(str(args.dir_log))
    ckpt_G = tf.train.Checkpoint(encoder=encoder, decoder=decoder)
    ckpt_manager = tf.train.CheckpointManager(ckpt_G,
                                              args.dir_checkpoint,
                                              max_to_keep=20)
    step = 0

    if args.dirs.checkpoint_G:
        _ckpt_manager = tf.train.CheckpointManager(ckpt_G,
                                                   args.dirs.checkpoint_G,
                                                   max_to_keep=1)
        ckpt_G.restore(_ckpt_manager.latest_checkpoint)
        print('checkpoint_G {} restored!!'.format(
            _ckpt_manager.latest_checkpoint))
        # cer = evaluate(feature_dev, dataset_dev, args.data.dev_size, encoder, decoder)
        # with writer.as_default():
        #     tf.summary.scalar("performance/cer", cer, step=step)

    start_time = datetime.now()
    num_processed = 0
    while step < 99999999:
        start = time()

        # supervise training
        uttids, x = next(iter_feature_train)
        trans = dataset_train.get_attrs('trans', uttids.numpy())
        loss_supervise = train_CTC_supervised(x, trans, encoder, decoder,
                                              optimizer)

        # unsupervise training
        text = next(iter_text)
        _, un_x = next(iter_feature_unsupervise)
        # loss_G = train_G(un_x, encoder, decoder, D, optimizer, args.model.D.max_label_len)
        loss_G = train_G(un_x, encoder, decoder, D, optimizer,
                         args.model.D.max_label_len)
        loss_D = train_D(un_x, text, encoder, decoder, D, optimizer_D,
                         args.lambda_gp, args.model.D.max_label_len)

        num_processed += len(un_x)
        progress = num_processed / args.data.untrain_size

        if step % 10 == 0:
            print(
                'loss_supervise: {:.3f}\tloss_G: {:.3f}\tloss_D: {:.3f}\tbatch: {}\tused: {:.3f}\t {:.3f}% step: {}'
                .format(loss_supervise, loss_G, loss_D, un_x.shape,
                        time() - start, progress * 100, step))
            with writer.as_default():
                tf.summary.scalar("costs/loss_supervise",
                                  loss_supervise,
                                  step=step)
        if step % args.dev_step == args.dev_step - 1:
            cer = evaluate(feature_dev, dataset_dev, args.data.dev_size,
                           encoder, decoder)
            with writer.as_default():
                tf.summary.scalar("performance/cer", cer, step=step)
        if step % args.decode_step == 0:
            monitor(dataset_dev[0], encoder, decoder)
        if step % args.save_step == 0:
            save_path = ckpt_manager.save(step)
            print('save model {}'.format(save_path))

        step += 1

    print('training duration: {:.2f}h'.format(
        (datetime.now() - start_time).total_seconds() / 3600))
def Train():
    dataset_train = ASR_align_DataSet(trans_file=args.dirs.train.trans,
                                      align_file=args.dirs.train.align,
                                      uttid2wav=args.dirs.train.wav_scp,
                                      feat_len_file=args.dirs.train.feat_len,
                                      args=args,
                                      _shuffle=True,
                                      transform=True)
    dataset_dev = ASR_align_DataSet(trans_file=args.dirs.dev.trans,
                                    align_file=args.dirs.dev.align,
                                    uttid2wav=args.dirs.dev.wav_scp,
                                    feat_len_file=args.dirs.dev.feat_len,
                                    args=args,
                                    _shuffle=False,
                                    transform=True)
    with tf.device("/cpu:0"):
        # wav data
        feature_train = TFData(dataset=dataset_train,
                               dir_save=args.dirs.train.tfdata,
                               args=args).read()
        feature_dev = TFData(dataset=dataset_dev,
                             dir_save=args.dirs.dev.tfdata,
                             args=args).read()
        if args.num_supervised:
            dataset_train_supervise = ASR_align_DataSet(
                trans_file=args.dirs.train_supervise.trans,
                align_file=args.dirs.train_supervise.align,
                uttid2wav=args.dirs.train_supervise.wav_scp,
                feat_len_file=args.dirs.train_supervise.feat_len,
                args=args,
                _shuffle=False,
                transform=True)
            feature_train_supervise = TFData(
                dataset=dataset_train_supervise,
                dir_save=args.dirs.train_supervise.tfdata,
                args=args).read()
            supervise_uttids, supervise_x = next(iter(feature_train_supervise.take(args.num_supervised).\
                padded_batch(args.num_supervised, ((), [None, args.dim_input]))))
            supervise_aligns = dataset_train_supervise.get_attrs(
                'align', supervise_uttids.numpy())
            # supervise_bounds = dataset_train_supervise.get_attrs('bounds', supervise_uttids.numpy())

        iter_feature_train = iter(
            feature_train.repeat().shuffle(500).padded_batch(
                args.batch_size,
                ((), [None, args.dim_input])).prefetch(buffer_size=5))
        feature_dev = feature_dev.padded_batch(args.batch_size,
                                               ((), [None, args.dim_input]))

    # create model paremeters
    model = PhoneClassifier(args)
    model.summary()
    optimizer_G = tf.keras.optimizers.Adam(args.opti.lr,
                                           beta_1=0.5,
                                           beta_2=0.9)

    writer = tf.summary.create_file_writer(str(args.dir_log))
    ckpt = tf.train.Checkpoint(model=model, optimizer_G=optimizer_G)
    ckpt_manager = tf.train.CheckpointManager(ckpt,
                                              args.dir_checkpoint,
                                              max_to_keep=20)
    step = 0

    # if a checkpoint exists, restore the latest checkpoint.
    if args.dirs.checkpoint:
        _ckpt_manager = tf.train.CheckpointManager(ckpt,
                                                   args.dirs.checkpoint,
                                                   max_to_keep=1)
        ckpt.restore(_ckpt_manager.latest_checkpoint)
        print('checkpoint {} restored!!'.format(
            _ckpt_manager.latest_checkpoint))
        step = int(_ckpt_manager.latest_checkpoint.split('-')[-1])

    start_time = datetime.now()

    while step < 99999999:
        start = time()

        if args.num_supervised:
            x = supervise_x
            loss_supervise = train_G_supervised(supervise_x, supervise_aligns,
                                                model, optimizer_G,
                                                args.dim_output)
            # loss_supervise, bounds_loss = train_G_bounds_supervised(
            #     x, supervise_bounds, supervise_aligns, model, optimizer_G, args.dim_output)
        else:
            uttids, x = next(iter_feature_train)
            aligns = dataset_train.get_attrs('align', uttids.numpy())
            # trans = dataset_train.get_attrs('trans', uttids.numpy())
            loss_supervise = train_G_supervised(x, aligns, model, optimizer_G,
                                                args.dim_output)
            # loss_supervise = train_G_TBTT_supervised(x, aligns, model, optimizer_G, args.dim_output)
            # bounds = dataset_train.get_attrs('bounds', uttids.numpy())
            # loss_supervise, bounds_loss = train_G_bounds_supervised(x, bounds, aligns, model, optimizer_G, args.dim_output)
            # loss_supervise = train_G_CTC_supervised(x, trans, model, optimizer_G)

        if step % 10 == 0:
            print('loss_supervise: {:.3f}\tbatch: {}\tused: {:.3f}\tstep: {}'.
                  format(loss_supervise, x.shape,
                         time() - start, step))
            # print('loss_supervise: {:.3f}\tloss_bounds: {:.3f}\tbatch: {}\tused: {:.3f}\tstep: {}'.format(
            #        loss_supervise, bounds_loss, x.shape, time()-start, step))
            with writer.as_default():
                tf.summary.scalar("costs/loss_supervise",
                                  loss_supervise,
                                  step=step)
        if step % args.dev_step == 0:
            fer, cer_0 = evaluate(feature_dev,
                                  dataset_dev,
                                  args.data.dev_size,
                                  model,
                                  beam_size=0,
                                  with_stamp=True)
            fer, cer = evaluate(feature_dev,
                                dataset_dev,
                                args.data.dev_size,
                                model,
                                beam_size=0,
                                with_stamp=False)
            with writer.as_default():
                tf.summary.scalar("performance/fer", fer, step=step)
                tf.summary.scalar("performance/cer_0", cer_0, step=step)
                tf.summary.scalar("performance/cer", cer, step=step)
        if step % args.decode_step == 0:
            monitor(dataset_dev[0], model)
        if step % args.save_step == 0:
            save_path = ckpt_manager.save(step)
            print('save model {}'.format(save_path))

        step += 1

    print('training duration: {:.2f}h'.format(
        (datetime.now() - start_time).total_seconds() / 3600))
def Train():
    with tf.device("/cpu:0"):
        dataset_train = ASR_align_ArkDataSet(scp_file=args.dirs.train.scp,
                                             trans_file=args.dirs.train.trans,
                                             align_file=None,
                                             feat_len_file=None,
                                             args=args,
                                             _shuffle=False,
                                             transform=True)

        feature_train = TFData(dataset=dataset_train,
                               dir_save=args.dirs.train.tfdata,
                               args=args).read(_shuffle=True, transform=True)
        bucket = tf.data.experimental.bucket_by_sequence_length(
            element_length_func=lambda uttid, x: tf.shape(x)[0],
            bucket_boundaries=args.list_bucket_boundaries,
            bucket_batch_sizes=args.list_batch_size,
            padded_shapes=((), [None, args.dim_input]))
        iter_feature_train = iter(
            feature_train.repeat().shuffle(10).apply(bucket).prefetch(
                buffer_size=5))

    # create model paremeters
    model = Transformer(args)
    model.summary()

    # learning_rate = CustomSchedule(args.model.G.d_model)
    # lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(
    #     0.0001,
    #     decay_steps=10000,
    #     decay_rate=0.5,
    #     staircase=True)
    optimizer = tf.keras.optimizers.Adam(0.000005,
                                         beta_1=0.5,
                                         beta_2=0.9,
                                         epsilon=1e-9)
    # optimizer = tf.keras.optimizers.SGD(0.1)

    ckpt = tf.train.Checkpoint(model=model)
    ckpt_manager = tf.train.CheckpointManager(ckpt,
                                              args.dir_checkpoint,
                                              max_to_keep=20)
    step = 0

    # if a checkpoint exists, restore the latest checkpoint.
    if args.dirs.checkpoint:
        _ckpt_manager = tf.train.CheckpointManager(ckpt,
                                                   args.dirs.checkpoint,
                                                   max_to_keep=1)
        ckpt.restore(_ckpt_manager.latest_checkpoint)
        print('checkpoint {} restored!!'.format(
            _ckpt_manager.latest_checkpoint))
        step = int(_ckpt_manager.latest_checkpoint.split('-')[-1])

    start_time = datetime.now()
    num_processed = 0
    # uttids, x = next(iter_feature_train)
    # trans_sos = dataset_train.get_attrs('trans_sos', uttids.numpy())
    # trans_eos = dataset_train.get_attrs('trans_eos', uttids.numpy())
    while step < 99999999:
        start = time()
        uttids, x = next(iter_feature_train)
        trans_sos = dataset_train.get_attrs('trans_sos', uttids.numpy())
        trans_eos = dataset_train.get_attrs('trans_eos', uttids.numpy())
        loss_supervise = train_step(x, trans_sos, trans_eos, model, optimizer)

        num_processed += len(x)
        progress = num_processed / args.data.train_size

        if step % 20 == 0:
            print('loss: {:.3f}\tbatch: {}\tused: {:.3f}\t {:.3f}% step: {}'.
                  format(loss_supervise, x.shape,
                         time() - start, progress * 100, step))
        # if step % args.save_step == 0:
        #     save_path = ckpt_manager.save(step)
        #     print('save model {}'.format(save_path))

        step += 1

    print('training duration: {:.2f}h'.format(
        (datetime.now() - start_time).total_seconds() / 3600))