|
- import argparse
- import json
- import os
-
- from mindspore import context, Tensor, ParameterTuple
- from mindspore.communication.management import init, get_rank, get_group_size
- from mindspore.communication import management as MultiDevice
- from mindspore.context import ParallelMode
- from mindspore.nn import TrainOneStepCell
- from mindspore.nn.optim import AdamWeightDecay
- from mindspore.ops.operations.math_ops import Mod
- from mindspore.train import Model
- from mindspore.train.callback import ModelCheckpoint, CheckpointConfig
- from mindspore.train.serialization import load_checkpoint, load_param_into_net
- from mindspore.train.loss_scale_manager import FixedLossScaleManager
-
- from src.callback import TimeMonitor, Monitor
- from src.config import train_config, symbols, encoder_kw, decoder_kw
- from src.dataset import create_dataset
- # from src.deepspeech2 import DeepSpeechModel, NetWithLossClass
- from src.model import Jasper, NetWithLossClass, init_weights
- from src.eval_callback import SaveCallback
- from src.lr_generator import get_lr
-
- parser = argparse.ArgumentParser(description='Jasper training')
- parser.add_argument('--pre_trained_model_path', type=str, default='', help='Pretrained checkpoint path')
- parser.add_argument('--is_distributed', action="store_true", default=False, help='Distributed training')
- parser.add_argument('--bidirectional', action="store_false", default=True, help='Use bidirectional RNN')
- parser.add_argument('--device_target', type=str, default="Ascend",
- help='Device target, support GPU and CPU, Default: GPU')
- args = parser.parse_args()
-
- if __name__ == '__main__':
-
- rank_id = 0
- group_size = 1
- config = train_config
- data_sink = False
- context.set_context(mode=context.GRAPH_MODE, device_target=args.device_target, save_graphs=False)
- # context.set_context(mode=context.PYNATIVE_MODE, device_target=args.device_target, save_graphs=False)
- if args.device_target == "GPU":
- context.set_context(enable_graph_kernel=False)
- if args.is_distributed:
- if args.device_target == "GPU":
- init()
- rank_id = get_rank()
- group_size = get_group_size()
- context.reset_auto_parallel_context()
- context.set_auto_parallel_context(device_num=get_group_size(), parallel_mode=ParallelMode.DATA_PARALLEL,
- gradients_mean=True)
- elif args.device_target == "Ascend":
- MultiDevice.init()
- context.reset_auto_parallel_context()
- context.set_auto_parallel_context(parallel_mode=ParallelMode.DATA_PARALLEL,
- device_num=MultiDevice.get_group_size(),
- gradients_mean=True)
- rank_id = MultiDevice.get_rank()
- group_size = MultiDevice.get_group_size()
- print(f"rank-{rank_id};Starting traning on multiple devices. |~ _ ~| |~ _ ~| |~ _ ~| |~ _ ~|")
-
- with open(config.DataConfig.labels_path) as label_file:
- labels = json.load(label_file)
-
- ds_train = create_dataset(data_dir=config.DataConfig.Data_dir,
- manifest_filepath=config.DataConfig.train_manifest,
- labels=symbols, batch_size=config.DataConfig.batch_size, train_mode=True,
- rank=rank_id, group_size=group_size)
- steps_size = ds_train.get_dataset_size()
-
- lr = get_lr(lr_init=config.OptimConfig.learning_rate, total_epochs=config.TrainingConfig.epochs,
- steps_per_epoch=steps_size)
- lr = Tensor(lr)
-
- # deepspeech_net = DeepSpeechModel(batch_size=config.DataConfig.batch_size,
- # rnn_hidden_size=config.ModelConfig.hidden_size,
- # nb_layers=config.ModelConfig.hidden_layers,
- # labels=labels,
- # rnn_type=config.ModelConfig.rnn_type,
- # audio_conf=config.DataConfig.SpectConfig,
- # bidirectional=True,
- # device_target=args.device_target)
- jasper_net = Jasper(encoder_kw=encoder_kw, decoder_kw=decoder_kw)
-
- loss_net = NetWithLossClass(jasper_net)
-
- init_weights(loss_net)
-
- weights = ParameterTuple(jasper_net.trainable_params())
-
- # optimizer = Adam(weights, learning_rate=config.OptimConfig.learning_rate, eps=config.OptimConfig.epsilon,
- # loss_scale=config.OptimConfig.loss_scale)
- optimizer = AdamWeightDecay(weights, learning_rate=lr, eps=config.OptimConfig.epsilon, weight_decay=1e-3)
- train_net = TrainOneStepCell(loss_net, optimizer)
- train_net.set_train(True)
- if args.pre_trained_model_path != '':
- param_dict = load_checkpoint(args.pre_trained_model_path)
- # load_param_into_net(train_net, param_dict)
- load_param_into_net(loss_net, param_dict)
- print('Successfully loading the pre-trained model')
-
- # model = Model(train_net)
- loss_scale = 128.0
- loss_scale = FixedLossScaleManager(loss_scale, drop_overflow_update=True)
- # model = TrainingWrapper(loss_net, optimizer, loss_scale)
-
- # model = Model(model)
- model = Model(loss_net, optimizer=optimizer, loss_scale_manager=loss_scale)
- # callback_list = [TimeMonitor(steps_size), Monitor(lr)]
- callback_list = [Monitor(lr)]
-
- if args.is_distributed:
- print('Distributed training.')
- config.CheckpointConfig.ckpt_path = os.path.join(config.CheckpointConfig.ckpt_path,
- 'ckpt_' + str(get_rank()) + '/')
- if rank_id == 0:
- # callback_update = SaveCallback(config.CheckpointConfig.ckpt_path)
- # callback_list += [callback_update]
- config_ck = CheckpointConfig(save_checkpoint_steps=1000,
- keep_checkpoint_max=config.CheckpointConfig.keep_checkpoint_max)
- ckpt_cb = ModelCheckpoint(prefix=config.CheckpointConfig.ckpt_file_name_prefix,
- directory=config.CheckpointConfig.ckpt_path, config=config_ck)
- callback_list.append(ckpt_cb)
- else:
- print('Standalone training.')
- config_ck = CheckpointConfig(save_checkpoint_steps=1000,
- keep_checkpoint_max=config.CheckpointConfig.keep_checkpoint_max)
- ckpt_cb = ModelCheckpoint(prefix=config.CheckpointConfig.ckpt_file_name_prefix,
- directory=config.CheckpointConfig.ckpt_path, config=config_ck)
-
- callback_list.append(ckpt_cb)
- print(callback_list)
- model.train(config.TrainingConfig.epochs, ds_train, callbacks=callback_list, dataset_sink_mode=data_sink)
|