|
- # Copyright 2021 Huawei Technologies Co., Ltd
- #
- # 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.
- # ============================================================================
-
- """train_criteo."""
-
- import argparse
- import json
- import os
- os.system("pip install -i https://pypi.tuna.tsinghua.edu.cn/simple librosa==0.9.2")
- os.system("pip install -i https://pypi.tuna.tsinghua.edu.cn/simple soundfile==0.11.0")
- os.system("pip install -i https://pypi.tuna.tsinghua.edu.cn/simple Levenshtein==0.20.9")
- from mindspore import context, Tensor, ParameterTuple
- from mindspore.communication.management import init, get_rank, get_group_size
- from mindspore.context import ParallelMode
- from mindspore.nn import TrainOneStepCell
- from mindspore.nn.optim import Adam
- from mindspore.train import Model
- from mindspore.train.callback import ModelCheckpoint, CheckpointConfig, LossMonitor, TimeMonitor
- from mindspore.train.serialization import load_checkpoint, load_param_into_net, save_checkpoint
-
- from src.config import train_config
- from src.dataset import create_dataset
- from src.deepspeech2 import DeepSpeechModel, NetWithLossClass
- from src.eval_callback import SaveCallback
- from src.lr_generator import get_lr
-
- parser = argparse.ArgumentParser(description='DeepSpeech2 training')
- parser.add_argument('--pre_trained_model_path', type=str, default='', help='Pretrained checkpoint path')
- parser.add_argument('--is_distributed', default=False, help='Distributed training')
- parser.add_argument('--run_modelarts', 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", choices=("GPU", "CPU", "Ascend"),
- help='Device target, support GPU and CPU, Default: GPU')
- parser.add_argument('--device_id', default=2, type=int, metavar='N', help='number of total epochs to run')
- parser.add_argument('--dataset_path', type=str, default="./data", help='Dataset path.')
- parser.add_argument('--data_url',
- help='path to training/inference dataset folder',
- default='./data')
-
- parser.add_argument('--train_url',
- help='model folder to save/load',
- default='./model')
-
- parser.add_argument('--result_url',
- help='folder to save inference results',
- default='./result')
-
- parser.add_argument('--device_num', type=int, default=1, help='Device num.')
- parser.add_argument('--epoch_size', type=int, default=70, help='Epoch size.')
- parser.add_argument('--checkpoint_path', type=str, default=None, help='CheckPoint file path.')
- args = parser.parse_args()
-
-
- if __name__ == '__main__':
-
- import os
- #print(os.listdir(args.data_url))
- ###新增代码
-
- rank_id = 0
- group_size = 1
- config = train_config
-
- ###新增代码s
-
- ###
- data_sink = (args.device_target != "CPU")
- context.set_context(mode=context.GRAPH_MODE, device_target=args.device_target, save_graphs=False)
- if args.device_target == "GPU":
- context.set_context(enable_graph_kernel=True)
- if args.is_distributed:
- init()
- #rank_id = get_rank()
- #group_size = get_group_size()
- rank_id = int(os.getenv('DEVICE_ID'))
- group_size = int(os.getenv('RANK_SIZE'))
-
- context.reset_auto_parallel_context()
- context.set_auto_parallel_context(device_num=get_group_size(), parallel_mode=ParallelMode.DATA_PARALLEL,
- gradients_mean=True, parameter_broadcast=True)
- if args.run_modelarts:
- import moxing as mox
- obs_data_url = args.data_url
- args.data_url = '/cache/data/'
- obs_train_url = args.train_url
- args.train_url = '/home/work/user-job-dir/outputs/model/'
- device_id = int(os.getenv('DEVICE_ID'))
- context.set_context(device_id=device_id)
- args.data_url = os.path.join(args.data_url, str(device_id))
- mox.file.copy_parallel(obs_data_url, args.data_url)
- print(os.listdir(args.data_url))
- config.DataConfig.labels_path = os.path.join(args.data_url, config.DataConfig.labels_path)
- config.DataConfig.train_manifest = os.path.join(args.data_url,config.DataConfig.train_manifest)
- config.CheckpointConfig.ckpt_path = args.data_url
- config.TrainingConfig.epochs = args.epoch_size
- else:
- if args.run_modelarts:
- device_id = int(os.getenv('DEVICE_ID'))
- context.set_context(device_id=device_id)
- import moxing as mox
- obs_data_url = args.data_url
- args.data_url = '/cache/data/'
- obs_train_url = args.train_url
- args.train_url = '/home/work/user-job-dir/outputs/model/'
-
- args.data_url = os.path.join(args.data_url, str(device_id))
- mox.file.copy_parallel(obs_data_url, args.data_url)
- print(os.listdir(args.data_url))
- config.DataConfig.labels_path = os.path.join(args.data_url, config.DataConfig.labels_path)
- config.DataConfig.train_manifest = os.path.join(args.data_url,config.DataConfig.train_manifest)
- config.CheckpointConfig.ckpt_path = args.data_url
- config.TrainingConfig.epochs = args.epoch_size
- elif args.device_target == 'Ascend':
- device_id = int(args.device_id)
- #device_id = int(os.getenv('DEVICE_ID'))
- context.set_context(device_id=device_id)
-
- with open(config.DataConfig.labels_path) as label_file:
- labels = json.load(label_file)
-
- ds_train = create_dataset(audio_conf=config.DataConfig.SpectConfig,
- manifest_filepath=config.DataConfig.train_manifest,
- labels=labels, normalize=True, train_mode=True,
- batch_size=config.DataConfig.batch_size, rank=rank_id,
- group_size=group_size,data_url=args.data_url,is_modelarts=args.run_modelarts)
- steps_size = ds_train.get_dataset_size()
- print("step_size:",steps_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)
-
- loss_net = NetWithLossClass(deepspeech_net, ascend=(args.device_target == 'Ascend'))
- weights = ParameterTuple(deepspeech_net.trainable_params())
-
- optimizer = Adam(weights, learning_rate=config.OptimConfig.learning_rate, eps=config.OptimConfig.epsilon,
- loss_scale=config.OptimConfig.loss_scale)
- 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)
- print('Successfully loading the pre-trained model')
-
- model = Model(train_net)
- callback_list = [TimeMonitor(steps_size), LossMonitor()]
-
- if args.is_distributed:
- config.CheckpointConfig.ckpt_path = os.path.join(config.CheckpointConfig.ckpt_path,
- 'ckpt_' + str(os.getenv('DEVICE_ID')) + '/')
- if args.run_modelarts:
- config.CheckpointConfig.ckpt_path = os.path.join(args.data_url,'ckpt')
- # if rank_id == 0:
-
- # callback_update = SaveCallback(config.CheckpointConfig.ckpt_path,args.data_url)
- # callback_list += [callback_update]
- config_ck = CheckpointConfig(save_checkpoint_steps=5,
- 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:
- config_ck = CheckpointConfig(save_checkpoint_steps=5,
- 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",callback_list)
- print("callback_list1",[TimeMonitor(steps_size), LossMonitor()])
- model.train(config.TrainingConfig.epochs, ds_train, callbacks=callback_list, dataset_sink_mode=data_sink)
-
- ###新增代码
- if args.run_modelarts:
-
- try:
- #output_path = "/home/work/user-job-dir/outputs/model/"
- #config.CheckpointConfig.ckpt_path = output_path
- #config.CheckpointConfig.ckpt_path = os.path.join(args.data_url,'ckpt')
-
- #os.mkdir(config.CheckpointConfig.ckpt_path)
- #print("ckpt_path",os.listdir(output_path))
-
- save_checkpoint(deepspeech_net,'{0}/deepspeech_epoch_{1}.ckpt'.format(config.CheckpointConfig.ckpt_path, config.TrainingConfig.epochs))
- print("ckpt_path1",os.listdir(config.CheckpointConfig.ckpt_path))
- mox.file.copy_parallel(config.CheckpointConfig.ckpt_path, obs_train_url)
-
- print("Successfully Upload {} to {}".format(config.CheckpointConfig.ckpt_path, obs_train_url))
- except Exception as e:
- print('moxing upload {} to {} failed: '.format(config.CheckpointConfig.ckpt_path, obs_train_url) + str(e))
- ###新增代码
|