|
- import os
- import json
- import librosa
- import moxing as mox
- import argparse
- from data_test1 import DatasetGenerator
- import mindspore.dataset as ds
- from mindspore import Model, load_checkpoint, load_param_into_net
- from mindspore import nn, context
- from mindspore.train.callback import LossMonitor, TimeMonitor, ModelCheckpoint, CheckpointConfig
- from network_define import WithLossCell
- from Loss_final1 import loss
- from model_rnn import Dual_RNN_model
- import time
-
- parser = argparse.ArgumentParser(
- description='Parameters for training Dual-Path-RNN')
-
- parser.add_argument('--in_dir', type=str, default=r"/home/work/user-job-dir/inputs/data/",
- help='Directory path of wsj0 including tr, cv and tt')
- parser.add_argument('--out_dir', type=str, default=r"/home/work/user-job-dir/inputs/data_json",
- help='Directory path to put output files')
- 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(
- '--device_target',
- type=str,
- default="Ascend",
- choices=['Ascend', 'GPU', 'CPU'],
- help='device where the code will be implemented (default: Ascend)')
-
- parser.add_argument('--train_dir', type=str, default=r"/home/work/user-job-dir/inputs/data_json/tr",
- help='directory including mix.json, s1.json and s2.json')
- parser.add_argument('--sample_rate', default=8000, type=int,
- help='Sample rate')
- parser.add_argument('--segment', default=4, type=float, # 取音频的长度,2s。#数据集语音长度要相同
- help='Segment length (seconds)')
- parser.add_argument('--batch_size', default=2, type=int, # 需要抛弃的音频长度
- help='Batch size')
-
- # Network architecture
- parser.add_argument('--in_channels', default=256, type=int,
- help='The number of expected features in the input')
- parser.add_argument('--out_channels', default=64, type=int,
- help='The number of features in the hidden state')
- parser.add_argument('--hidden_channels', default=128, type=int,
- help='The hidden size of RNN')
- parser.add_argument('--kernel_size', default=2, type=int,
- help='Encoder and Decoder Kernel size')
- parser.add_argument('--rnn_type', default='LSTM', type=str,
- help='RNN, LSTM, GRU')
- parser.add_argument('--norm', default='gln', type=str,
- help='gln = "Global Norm", cln = "Cumulative Norm", ln = "Layer Norm"')
- parser.add_argument('--dropout', default=0.0, type=float,
- help='dropout')
- parser.add_argument('--num_layers', default=4, type=int,
- help='Number of Dual-Path-Block')
- parser.add_argument('--K', default=250, type=int,
- help='The length of chunk')
- parser.add_argument('--num_spks', default=2, type=int,
- help='The number of speakers')
-
- # optimizer
- parser.add_argument('--lr', default=1e-3, type=float,
- help='Init learning rate')
- parser.add_argument('--l2', default=1e-5, type=float,
- help='weight decay (L2 penalty)')
-
- # save and load model
- parser.add_argument('--save_checkpoint_path', default=r"./checkpoint",
- help='Location to save epoch models')
-
- def preprocess_one_dir(in_dir, out_dir, out_filename, sample_rate=8000):
- file_infos = []
- in_dir = os.path.abspath(in_dir)
- wav_list = os.listdir(in_dir)
- for wav_file in wav_list:
- if not wav_file.endswith('.wav'):
- continue
- wav_path = os.path.join(in_dir, wav_file)
- samples, _ = librosa.load(wav_path, sr=sample_rate)
- file_infos.append((wav_path, len(samples)))
- if not os.path.exists(out_dir):
- os.makedirs(out_dir)
- with open(os.path.join(out_dir, out_filename + '.json'), 'w') as f:
- json.dump(file_infos, f, indent=4)
-
- def preprocess(args):
- for data_type in ['tr']:
- for speaker in ['mix', 's1', 's2']:
- preprocess_one_dir(os.path.join(args.in_dir, data_type, speaker),
- os.path.join(args.out_dir, data_type),
- speaker,
- sample_rate=args.sample_rate)
- print("preprocess done")
-
- def main(args):
- # context.set_context(mode=context.GRAPH_MODE, device_target=args.device_target, save_graphs=True)
- context.set_context(mode=context.GRAPH_MODE, device_target=args.device_target)
-
- # 在训练环境中定义data_url和train_url,并把数据从obs拷贝到相应的固定路径
- obs_data_url = args.data_url
- args.data_url = '/home/work/user-job-dir/inputs/data/'
- obs_train_url = args.train_url
- args.train_url = '/home/work/user-job-dir/outputs/model/'
- try:
- mox.file.copy_parallel(obs_data_url, args.data_url)
- print("Successfully Download {} to {}".format(obs_data_url, args.data_url))
- except Exception as e:
- print('moxing download {} to {} failed: '.format(
- obs_data_url, args.data_url) + str(e))
-
- print("start preprocess ....")
- preprocess(args)
- args.save_checkpoint_path = args.train_url
-
- # build dataloader
- tr_dataset = DatasetGenerator(args.train_dir, args.batch_size,
- sample_rate=args.sample_rate, segment=args.segment)
- tr_loader = ds.GeneratorDataset(tr_dataset, ["mixture", "lens", "sources"], shuffle=True)
-
- tr_loader = tr_loader.batch(1)
- num_steps = tr_loader.get_dataset_size()
-
- # build model
- net = Dual_RNN_model(args.in_channels, args.out_channels, args.hidden_channels,
- bidirectional=True, norm=args.norm, num_layers=args.num_layers, dropout=args.dropout, K=args.K)
-
- print(net)
- net.set_train()
- # build optimizer
- optimizier = nn.Adam(net.get_parameters(), learning_rate=args.lr, weight_decay=args.l2)
- my_loss = loss()
- net_with_loss = WithLossCell(net, my_loss)
- model = Model(net_with_loss, optimizer=optimizier)
-
- loss_cb = LossMonitor(1)
- time_cb = TimeMonitor()
- cb = [time_cb, loss_cb]
-
- config_ck = CheckpointConfig(save_checkpoint_steps=100, keep_checkpoint_max=5)
- ckpt_cb = ModelCheckpoint(prefix='dual',
- directory=args.save_checkpoint_path,
- config=config_ck)
- cb += [ckpt_cb]
-
- #开始训练
- print("============== Starting Training ==============")
- model.train(epoch=1, train_dataset=tr_loader, callbacks=cb, dataset_sink_mode=False)
-
- ######################## 将输出的模型拷贝到obs(固定写法) ########################
- try:
- mox.file.copy_parallel(args.train_url, obs_train_url)
- print("Successfully Upload {} to {}".format(args.train_url,
- obs_train_url))
- except Exception as e:
- print('moxing upload {} to {} failed: '.format(args.train_url,
- obs_train_url) + str(e))
-
- if __name__ == '__main__':
- args = parser.parse_args()
- print(args)
- main(args)
|