Hui Zhang
3 years ago
parent
91bc5959a9
commit
3843372958
@ -0,0 +1,219 @@
|
|||||||
|
# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
|
||||||
|
#
|
||||||
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
# you may not use this file except in compliance with the License.
|
||||||
|
# You may obtain a copy of the License at
|
||||||
|
#
|
||||||
|
# http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
#
|
||||||
|
# Unless required by applicable law or agreed to in writing, software
|
||||||
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
# See the License for the specific language governing permissions and
|
||||||
|
# limitations under the License.
|
||||||
|
"""Contains U2 model."""
|
||||||
|
import paddle
|
||||||
|
from paddle import distributed as dist
|
||||||
|
from paddle.io import DataLoader
|
||||||
|
|
||||||
|
from deepspeech.io.collator import SpeechCollator
|
||||||
|
from deepspeech.io.dataset import ManifestDataset
|
||||||
|
from deepspeech.io.sampler import SortagradBatchSampler
|
||||||
|
from deepspeech.io.sampler import SortagradDistributedBatchSampler
|
||||||
|
from deepspeech.models.u2 import U2Evaluator
|
||||||
|
from deepspeech.models.u2 import U2Model
|
||||||
|
from deepspeech.models.u2 import U2Updater
|
||||||
|
from deepspeech.training.extensions.snapshot import Snapshot
|
||||||
|
from deepspeech.training.extensions.visualizer import VisualDL
|
||||||
|
from deepspeech.training.optimizer import OptimizerFactory
|
||||||
|
from deepspeech.training.scheduler import LRSchedulerFactory
|
||||||
|
from deepspeech.training.timer import Timer
|
||||||
|
from deepspeech.training.trainer import Trainer
|
||||||
|
from deepspeech.training.updaters.trainer import Trainer as NewTrainer
|
||||||
|
from deepspeech.utils import layer_tools
|
||||||
|
from deepspeech.utils.log import Log
|
||||||
|
|
||||||
|
logger = Log(__name__).getlog()
|
||||||
|
|
||||||
|
|
||||||
|
class U2Trainer(Trainer):
|
||||||
|
def __init__(self, config, args):
|
||||||
|
super().__init__(config, args)
|
||||||
|
|
||||||
|
def setup_dataloader(self):
|
||||||
|
config = self.config.clone()
|
||||||
|
config.defrost()
|
||||||
|
config.collator.keep_transcription_text = False
|
||||||
|
|
||||||
|
# train/valid dataset, return token ids
|
||||||
|
config.data.manifest = config.data.train_manifest
|
||||||
|
train_dataset = ManifestDataset.from_config(config)
|
||||||
|
|
||||||
|
config.data.manifest = config.data.dev_manifest
|
||||||
|
dev_dataset = ManifestDataset.from_config(config)
|
||||||
|
|
||||||
|
collate_fn_train = SpeechCollator.from_config(config)
|
||||||
|
|
||||||
|
config.collator.augmentation_config = ""
|
||||||
|
collate_fn_dev = SpeechCollator.from_config(config)
|
||||||
|
|
||||||
|
if self.parallel:
|
||||||
|
batch_sampler = SortagradDistributedBatchSampler(
|
||||||
|
train_dataset,
|
||||||
|
batch_size=config.collator.batch_size,
|
||||||
|
num_replicas=None,
|
||||||
|
rank=None,
|
||||||
|
shuffle=True,
|
||||||
|
drop_last=True,
|
||||||
|
sortagrad=config.collator.sortagrad,
|
||||||
|
shuffle_method=config.collator.shuffle_method)
|
||||||
|
else:
|
||||||
|
batch_sampler = SortagradBatchSampler(
|
||||||
|
train_dataset,
|
||||||
|
shuffle=True,
|
||||||
|
batch_size=config.collator.batch_size,
|
||||||
|
drop_last=True,
|
||||||
|
sortagrad=config.collator.sortagrad,
|
||||||
|
shuffle_method=config.collator.shuffle_method)
|
||||||
|
self.train_loader = DataLoader(
|
||||||
|
train_dataset,
|
||||||
|
batch_sampler=batch_sampler,
|
||||||
|
collate_fn=collate_fn_train,
|
||||||
|
num_workers=config.collator.num_workers, )
|
||||||
|
self.valid_loader = DataLoader(
|
||||||
|
dev_dataset,
|
||||||
|
batch_size=config.collator.batch_size,
|
||||||
|
shuffle=False,
|
||||||
|
drop_last=False,
|
||||||
|
collate_fn=collate_fn_dev)
|
||||||
|
|
||||||
|
# test dataset, return raw text
|
||||||
|
config.data.manifest = config.data.test_manifest
|
||||||
|
# filter test examples, will cause less examples, but no mismatch with training
|
||||||
|
# and can use large batch size , save training time, so filter test egs now.
|
||||||
|
config.data.min_input_len = 0.0 # second
|
||||||
|
config.data.max_input_len = float('inf') # second
|
||||||
|
config.data.min_output_len = 0.0 # tokens
|
||||||
|
config.data.max_output_len = float('inf') # tokens
|
||||||
|
config.data.min_output_input_ratio = 0.00
|
||||||
|
config.data.max_output_input_ratio = float('inf')
|
||||||
|
|
||||||
|
test_dataset = ManifestDataset.from_config(config)
|
||||||
|
# return text ord id
|
||||||
|
config.collator.keep_transcription_text = True
|
||||||
|
config.collator.augmentation_config = ""
|
||||||
|
self.test_loader = DataLoader(
|
||||||
|
test_dataset,
|
||||||
|
batch_size=config.decoding.batch_size,
|
||||||
|
shuffle=False,
|
||||||
|
drop_last=False,
|
||||||
|
collate_fn=SpeechCollator.from_config(config))
|
||||||
|
# return text token id
|
||||||
|
config.collator.keep_transcription_text = False
|
||||||
|
self.align_loader = DataLoader(
|
||||||
|
test_dataset,
|
||||||
|
batch_size=config.decoding.batch_size,
|
||||||
|
shuffle=False,
|
||||||
|
drop_last=False,
|
||||||
|
collate_fn=SpeechCollator.from_config(config))
|
||||||
|
logger.info("Setup train/valid/test/align Dataloader!")
|
||||||
|
|
||||||
|
def setup_model(self):
|
||||||
|
config = self.config
|
||||||
|
model_conf = config.model
|
||||||
|
model_conf.defrost()
|
||||||
|
model_conf.input_dim = self.train_loader.collate_fn.feature_size
|
||||||
|
model_conf.output_dim = self.train_loader.collate_fn.vocab_size
|
||||||
|
model_conf.freeze()
|
||||||
|
model = U2Model.from_config(model_conf)
|
||||||
|
|
||||||
|
if self.parallel:
|
||||||
|
model = paddle.DataParallel(model)
|
||||||
|
|
||||||
|
model.train()
|
||||||
|
logger.info(f"{model}")
|
||||||
|
layer_tools.print_params(model, logger.info)
|
||||||
|
|
||||||
|
train_config = config.training
|
||||||
|
optim_type = train_config.optim
|
||||||
|
optim_conf = train_config.optim_conf
|
||||||
|
scheduler_type = train_config.scheduler
|
||||||
|
scheduler_conf = train_config.scheduler_conf
|
||||||
|
|
||||||
|
scheduler_args = {
|
||||||
|
"learning_rate": optim_conf.lr,
|
||||||
|
"verbose": False,
|
||||||
|
"warmup_steps": scheduler_conf.warmup_steps,
|
||||||
|
"gamma": scheduler_conf.lr_decay,
|
||||||
|
"d_model": model_conf.encoder_conf.output_size,
|
||||||
|
}
|
||||||
|
lr_scheduler = LRSchedulerFactory.from_args(scheduler_type,
|
||||||
|
scheduler_args)
|
||||||
|
|
||||||
|
def optimizer_args(
|
||||||
|
config,
|
||||||
|
parameters,
|
||||||
|
lr_scheduler=None, ):
|
||||||
|
train_config = config.training
|
||||||
|
optim_type = train_config.optim
|
||||||
|
optim_conf = train_config.optim_conf
|
||||||
|
scheduler_type = train_config.scheduler
|
||||||
|
scheduler_conf = train_config.scheduler_conf
|
||||||
|
return {
|
||||||
|
"grad_clip": train_config.global_grad_clip,
|
||||||
|
"weight_decay": optim_conf.weight_decay,
|
||||||
|
"learning_rate": lr_scheduler
|
||||||
|
if lr_scheduler else optim_conf.lr,
|
||||||
|
"parameters": parameters,
|
||||||
|
"epsilon": 1e-9 if optim_type == 'noam' else None,
|
||||||
|
"beta1": 0.9 if optim_type == 'noam' else None,
|
||||||
|
"beat2": 0.98 if optim_type == 'noam' else None,
|
||||||
|
}
|
||||||
|
|
||||||
|
optimzer_args = optimizer_args(config, model.parameters(), lr_scheduler)
|
||||||
|
optimizer = OptimizerFactory.from_args(optim_type, optimzer_args)
|
||||||
|
|
||||||
|
self.model = model
|
||||||
|
self.optimizer = optimizer
|
||||||
|
self.lr_scheduler = lr_scheduler
|
||||||
|
logger.info("Setup model/optimizer/lr_scheduler!")
|
||||||
|
|
||||||
|
def setup_updater(self):
|
||||||
|
output_dir = self.output_dir
|
||||||
|
config = self.config.training
|
||||||
|
|
||||||
|
updater = U2Updater(
|
||||||
|
model=self.model,
|
||||||
|
optimizer=self.optimizer,
|
||||||
|
scheduler=self.lr_scheduler,
|
||||||
|
dataloader=self.train_loader,
|
||||||
|
output_dir=output_dir,
|
||||||
|
accum_grad=config.accum_grad)
|
||||||
|
|
||||||
|
trainer = NewTrainer(updater, (config.n_epoch, 'epoch'), output_dir)
|
||||||
|
|
||||||
|
evaluator = U2Evaluator(self.model, self.valid_loader)
|
||||||
|
|
||||||
|
trainer.extend(evaluator, trigger=(1, "epoch"))
|
||||||
|
|
||||||
|
if dist.get_rank() == 0:
|
||||||
|
trainer.extend(VisualDL(output_dir), trigger=(1, "iteration"))
|
||||||
|
num_snapshots = config.checkpoint.kbest_n
|
||||||
|
trainer.extend(
|
||||||
|
Snapshot(
|
||||||
|
mode='kbest',
|
||||||
|
max_size=num_snapshots,
|
||||||
|
indicator='VALID/LOSS',
|
||||||
|
less_better=True),
|
||||||
|
trigger=(1, 'epoch'))
|
||||||
|
# print(trainer.extensions)
|
||||||
|
# trainer.run()
|
||||||
|
self.trainer = trainer
|
||||||
|
|
||||||
|
def run(self):
|
||||||
|
"""The routine of the experiment after setup. This method is intended
|
||||||
|
to be used by the user.
|
||||||
|
"""
|
||||||
|
self.setup_updater()
|
||||||
|
with Timer("Training Done: {}"):
|
||||||
|
self.trainer.run()
|
||||||
@ -0,0 +1,19 @@
|
|||||||
|
# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
|
||||||
|
#
|
||||||
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
# you may not use this file except in compliance with the License.
|
||||||
|
# You may obtain a copy of the License at
|
||||||
|
#
|
||||||
|
# http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
#
|
||||||
|
# Unless required by applicable law or agreed to in writing, software
|
||||||
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
# See the License for the specific language governing permissions and
|
||||||
|
# limitations under the License.
|
||||||
|
from .u2 import U2InferModel
|
||||||
|
from .u2 import U2Model
|
||||||
|
from .updater import U2Evaluator
|
||||||
|
from .updater import U2Updater
|
||||||
|
|
||||||
|
__all__ = ["U2Model", "U2InferModel", "U2Evaluator", "U2Updater"]
|
||||||
@ -0,0 +1,149 @@
|
|||||||
|
# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
|
||||||
|
#
|
||||||
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
# you may not use this file except in compliance with the License.
|
||||||
|
# You may obtain a copy of the License at
|
||||||
|
#
|
||||||
|
# http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
#
|
||||||
|
# Unless required by applicable law or agreed to in writing, software
|
||||||
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
# See the License for the specific language governing permissions and
|
||||||
|
# limitations under the License.
|
||||||
|
from contextlib import nullcontext
|
||||||
|
|
||||||
|
import paddle
|
||||||
|
from paddle import distributed as dist
|
||||||
|
|
||||||
|
from deepspeech.training.extensions.evaluator import StandardEvaluator
|
||||||
|
from deepspeech.training.reporter import report
|
||||||
|
from deepspeech.training.timer import Timer
|
||||||
|
from deepspeech.training.updaters.standard_updater import StandardUpdater
|
||||||
|
from deepspeech.utils import layer_tools
|
||||||
|
from deepspeech.utils.log import Log
|
||||||
|
|
||||||
|
logger = Log(__name__).getlog()
|
||||||
|
|
||||||
|
|
||||||
|
class U2Evaluator(StandardEvaluator):
|
||||||
|
def __init__(self, model, dataloader):
|
||||||
|
super().__init__(model, dataloader)
|
||||||
|
self.msg = ""
|
||||||
|
self.num_seen_utts = 0
|
||||||
|
self.total_loss = 0.0
|
||||||
|
|
||||||
|
def evaluate_core(self, batch):
|
||||||
|
self.msg = "Valid: Rank: {}, ".format(dist.get_rank())
|
||||||
|
losses_dict = {}
|
||||||
|
|
||||||
|
loss, attention_loss, ctc_loss = self.model(*batch[1:])
|
||||||
|
if paddle.isfinite(loss):
|
||||||
|
num_utts = batch[1].shape[0]
|
||||||
|
self.num_seen_utts += num_utts
|
||||||
|
self.total_loss += float(loss) * num_utts
|
||||||
|
|
||||||
|
losses_dict['loss'] = float(loss)
|
||||||
|
if attention_loss:
|
||||||
|
losses_dict['att_loss'] = float(attention_loss)
|
||||||
|
if ctc_loss:
|
||||||
|
losses_dict['ctc_loss'] = float(ctc_loss)
|
||||||
|
|
||||||
|
for k, v in losses_dict.items():
|
||||||
|
report("eval/" + k, v)
|
||||||
|
|
||||||
|
self.msg += ', '.join('{}: {:>.6f}'.format(k, v)
|
||||||
|
for k, v in losses_dict.items())
|
||||||
|
logger.info(self.msg)
|
||||||
|
return self.total_loss, self.num_seen_utts
|
||||||
|
|
||||||
|
|
||||||
|
class U2Updater(StandardUpdater):
|
||||||
|
def __init__(self,
|
||||||
|
model,
|
||||||
|
optimizer,
|
||||||
|
scheduler,
|
||||||
|
dataloader,
|
||||||
|
init_state=None,
|
||||||
|
accum_grad=1,
|
||||||
|
**kwargs):
|
||||||
|
super().__init__(
|
||||||
|
model, optimizer, scheduler, dataloader, init_state=init_state)
|
||||||
|
self.accum_grad = accum_grad
|
||||||
|
self.forward_count = 0
|
||||||
|
self.msg = ""
|
||||||
|
|
||||||
|
def update_core(self, batch):
|
||||||
|
"""One Step
|
||||||
|
|
||||||
|
Args:
|
||||||
|
batch (List[Object]): utts, xs, xlens, ys, ylens
|
||||||
|
"""
|
||||||
|
losses_dict = {}
|
||||||
|
self.msg = "Rank: {}, ".format(dist.get_rank())
|
||||||
|
|
||||||
|
# forward
|
||||||
|
batch_size = batch[1].shape[0]
|
||||||
|
loss, attention_loss, ctc_loss = self.model(*batch[1:])
|
||||||
|
# loss div by `batch_size * accum_grad`
|
||||||
|
loss /= self.accum_grad
|
||||||
|
|
||||||
|
# loss backward
|
||||||
|
if (self.forward_count + 1) != self.accum_grad:
|
||||||
|
# Disable gradient synchronizations across DDP processes.
|
||||||
|
# Within this context, gradients will be accumulated on module
|
||||||
|
# variables, which will later be synchronized.
|
||||||
|
context = self.model.no_sync
|
||||||
|
else:
|
||||||
|
# Used for single gpu training and DDP gradient synchronization
|
||||||
|
# processes.
|
||||||
|
context = nullcontext
|
||||||
|
|
||||||
|
with context():
|
||||||
|
loss.backward()
|
||||||
|
layer_tools.print_grads(self.model, print_func=None)
|
||||||
|
|
||||||
|
# loss info
|
||||||
|
losses_dict['loss'] = float(loss) * self.accum_grad
|
||||||
|
if attention_loss:
|
||||||
|
losses_dict['att_loss'] = float(attention_loss)
|
||||||
|
if ctc_loss:
|
||||||
|
losses_dict['ctc_loss'] = float(ctc_loss)
|
||||||
|
# report loss
|
||||||
|
for k, v in losses_dict.items():
|
||||||
|
report("train/" + k, v)
|
||||||
|
# loss msg
|
||||||
|
self.msg += "batch size: {}, ".format(batch_size)
|
||||||
|
self.msg += "accum: {}, ".format(self.accum_grad)
|
||||||
|
self.msg += ', '.join('{}: {:>.6f}'.format(k, v)
|
||||||
|
for k, v in losses_dict.items())
|
||||||
|
|
||||||
|
# Truncate the graph
|
||||||
|
loss.detach()
|
||||||
|
|
||||||
|
# update parameters
|
||||||
|
self.forward_count += 1
|
||||||
|
if self.forward_count != self.accum_grad:
|
||||||
|
return
|
||||||
|
self.forward_count = 0
|
||||||
|
|
||||||
|
self.optimizer.step()
|
||||||
|
self.optimizer.clear_grad()
|
||||||
|
self.scheduler.step()
|
||||||
|
|
||||||
|
def update(self):
|
||||||
|
# model is default in train mode
|
||||||
|
|
||||||
|
# training for a step is implemented here
|
||||||
|
with Timer("data time cost:{}"):
|
||||||
|
batch = self.read_batch()
|
||||||
|
with Timer("step time cost:{}"):
|
||||||
|
self.update_core(batch)
|
||||||
|
|
||||||
|
# #iterations with accum_grad > 1
|
||||||
|
# Ref.: https://github.com/espnet/espnet/issues/777
|
||||||
|
if self.forward_count == 0:
|
||||||
|
self.state.iteration += 1
|
||||||
|
if self.updates_per_epoch is not None:
|
||||||
|
if self.state.iteration % self.updates_per_epoch == 0:
|
||||||
|
self.state.epoch += 1
|
||||||
Loading…
Reference in new issue