r"""Train, evaluate, and save checkpoints. Args: workspace: str, directory of workspace gpus: int, number of GPUs to train config_yaml: str
(args)
| 137 | |
| 138 | |
| 139 | def train(args) -> NoReturn: |
| 140 | r"""Train, evaluate, and save checkpoints. |
| 141 | |
| 142 | Args: |
| 143 | workspace: str, directory of workspace |
| 144 | gpus: int, number of GPUs to train |
| 145 | config_yaml: str |
| 146 | """ |
| 147 | |
| 148 | # arguments & parameters |
| 149 | workspace = args.workspace |
| 150 | config_yaml = args.config_yaml |
| 151 | filename = args.filename |
| 152 | |
| 153 | devices_num = torch.cuda.device_count() |
| 154 | # Read config file. |
| 155 | configs = parse_yaml(config_yaml) |
| 156 | |
| 157 | # Configuration of data |
| 158 | max_mix_num = configs['data']['max_mix_num'] |
| 159 | sampling_rate = configs['data']['sampling_rate'] |
| 160 | lower_db = configs['data']['loudness_norm']['lower_db'] |
| 161 | higher_db = configs['data']['loudness_norm']['higher_db'] |
| 162 | |
| 163 | # Configuration of the separation model |
| 164 | query_net = configs['model']['query_net'] |
| 165 | model_type = configs['model']['model_type'] |
| 166 | input_channels = configs['model']['input_channels'] |
| 167 | output_channels = configs['model']['output_channels'] |
| 168 | condition_size = configs['model']['condition_size'] |
| 169 | use_text_ratio = configs['model']['use_text_ratio'] |
| 170 | |
| 171 | # Configuration of the trainer |
| 172 | num_nodes = configs['train']['num_nodes'] |
| 173 | batch_size = configs['train']['batch_size_per_device'] |
| 174 | sync_batchnorm = configs['train']['sync_batchnorm'] |
| 175 | num_workers = configs['train']['num_workers'] |
| 176 | loss_type = configs['train']['loss_type'] |
| 177 | optimizer_type = configs["train"]["optimizer"]["optimizer_type"] |
| 178 | learning_rate = float(configs['train']["optimizer"]['learning_rate']) |
| 179 | lr_lambda_type = configs['train']["optimizer"]['lr_lambda_type'] |
| 180 | warm_up_steps = configs['train']["optimizer"]['warm_up_steps'] |
| 181 | reduce_lr_steps = configs['train']["optimizer"]['reduce_lr_steps'] |
| 182 | save_step_frequency = configs['train']['save_step_frequency'] |
| 183 | resume_checkpoint_path = args.resume_checkpoint_path |
| 184 | if resume_checkpoint_path == "": |
| 185 | resume_checkpoint_path = None |
| 186 | else: |
| 187 | logging.info(f'Finetuning AudioSep with checkpoint [{resume_checkpoint_path}]') |
| 188 | |
| 189 | # Get directories and paths |
| 190 | checkpoints_dir, logs_dir, tf_logs_dir, statistics_path = get_dirs( |
| 191 | workspace, filename, config_yaml, devices_num, |
| 192 | ) |
| 193 | |
| 194 | logging.info(configs) |
| 195 | |
| 196 | # data module |
no test coverage detected