def main(): # Setup the log log_filename = config.LOG_FILENAME log = setup_log(log_filename) # Setup datasets instruments = config.INSTRUMENTS musdbwav_path = config.MUSDB18_WAV_PATH # sample_time_frames = config.SAMPLE_TIME_FRAMES audio_samples_per_chunk = config.AUDIO_SAMPLES_PER_CHUNK batch_size = config.BATCH_SIZE prefetch_factor = config.PREFETCH_FACTOR train_dataset = TrainMUSDB18Dataset(musdbwav_path, instruments, sample_length=audio_samples_per_chunk) validation_dataset = TestMUSDB18Dataset( musdbwav_path, instruments, sample_length=audio_samples_per_chunk, subset_split='valid') # Setup model model_config_path = config.MODEL_CONFIG_YAML_PATH model_config_name = config.MODEL_CONFIGURATION with open(model_config_path, 'r') as file: model_configurations = yaml.load(file, Loader=yaml.FullLoader) train_model_config = model_configurations[model_config_name] train_model_class = train_model_config.pop('class') model = eval(train_model_class)(**train_model_config) # Setup trainer checkpoint_folder_path = config.CHECKPOINT_FOLDER_PATH epochs = config.EPOCHS checkpoint_frequency = config.CHECKPOINT_FREQUENCY logging_frequency = config.LOGGING_FREQUENCY optimizer_class = get_optimizer_class(config.OPTIMIZER) optimizer_params = config.OPTIMIZER_PARAMS optimizer = optimizer_class(model.parameters(), **optimizer_params) lr_scheduler_class = get_lr_scheduler_class(config.LR_SCHEDULER) lr_scheduler_params = config.LR_SCHEDULER_PARAMS lr_scheduler = lr_scheduler_class(optimizer, **lr_scheduler_params) loss_function = get_loss_function(config.LOSS_FUNCTION) gpu = config.GPU gpu_device = config.GPU_DEVICE device = get_device(gpu, gpu_device) if not os.path.isdir(checkpoint_folder_path): os.mkdir(checkpoint_folder_path) # Setup the spectrogram spectrogram_type = config.SPECTROGRAM_TYPE n_fft = config.N_FFT hop_length = config.HOP_LENGTH window = config.WINDOW window_length = config.WINDOW_LENGTH spectrogramer = Spectrogramer(spectrogram_type, n_fft, hop_length, window, window_length, device) # Initialize traininer trainer = Trainer(model, spectrogramer, optimizer, loss_function, lr_scheduler, train_dataset, validation_dataset, log, checkpoint_folder_path, epochs, logging_frequency, checkpoint_frequency, batch_size, prefetch_factor, instruments, train_model_class, device) # Start trainer/evaluation trainer.train()
parser.add_argument('--load_order', type=str, default='-') parser.add_argument('--maml_lr', default=0.1, type=float) parser.add_argument('--maml_epoch', default=50, type=int) parser.add_argument('--mnemonics_images_per_class_per_step', default=1, type=int) parser.add_argument('--mnemonics_steps', default=20, type=int) parser.add_argument('--mnemonics_epochs', default=5, type=int) parser.add_argument('--mnemonics_lr', type=float, default=0.01) parser.add_argument('--mnemonics_decay_factor', type=float, default=0.5) parser.add_argument('--mnemonics_outer_lr', type=float, default=1e-6) parser.add_argument('--mnemonics_total_epochs', type=int, default=10) parser.add_argument('--mnemonics_decay_epochs', type=int, default=40) the_args = parser.parse_args() assert(the_args.nb_cl_fg % the_args.nb_cl == 0) assert(the_args.nb_cl_fg >= the_args.nb_cl) print(the_args) np.random.seed(the_args.random_seed) os.environ['CUDA_VISIBLE_DEVICES'] = the_args.gpu print('Using gpu:', the_args.gpu) occupy_memory(the_args.gpu) print('Occupy GPU memory in advance') trainer = Trainer(the_args) trainer.train()