Spaces:
Running
Running
| # Copyright (c) Meta Platforms, Inc. and affiliates. | |
| # All rights reserved. | |
| # | |
| # This source code is licensed under the license found in the | |
| # LICENSE file in the root directory of this source tree. | |
| """ | |
| dora grid watermarking.kbits --clear | |
| """ | |
| import os | |
| from audiocraft.environment import AudioCraftEnvironment | |
| from ._explorers import WatermarkingMbExplorer | |
| def explorer(launcher): | |
| partitions = AudioCraftEnvironment.get_slurm_partitions(['team', 'global']) | |
| launcher.slurm_( | |
| gpus=8, | |
| partition=partitions, | |
| constraint="volta32gb", | |
| ) | |
| launcher.bind_( | |
| { | |
| "solver": "watermark/robustness", | |
| "dset": os.getenv("AUDIOCRAFT_DSET", "audio/example"), | |
| "dataset.batch_size": 16, | |
| # optim | |
| "optim.epochs": 300, | |
| "schedule": { | |
| "lr_scheduler": "cosine", | |
| "cosine": { | |
| "warmup": 4000, | |
| "lr_min_ratio": 0.0, | |
| "cycle_length": 1.0, | |
| }, | |
| }, | |
| # crop and padding | |
| "crop": { | |
| "prob": 0.4, | |
| "shuffle_prob": 0.2, | |
| "pad_prob": 0.2, | |
| "size": 0.5, | |
| "max_n_windows": 5, | |
| }, | |
| # augmentations | |
| "select_aug_mode": 'use_eval', | |
| "aug_weights.updownresample": 0.1, | |
| "aug_weights.speed": 0.1, | |
| "aug_weights.echo": 0.1, | |
| "aug_weights.pink_noise": 0.1, | |
| "aug_weights.lowpass_filter": 0.1, | |
| "aug_weights.highpass_filter": 0.1, | |
| "aug_weights.bandpass_filter": 0.1, | |
| "aug_weights.smooth": 0.1, | |
| "aug_weights.boost_audio": 0.1, | |
| "aug_weights.duck_audio": 0.1, | |
| "aug_weights.mp3_compression": 0.1, | |
| "aug_weights.encodec": 0.1, | |
| "aug_weights.identity": 1.0, | |
| # multi-bit | |
| "audioseal.nbits": 16, | |
| "detector.output_dim": 32, | |
| "wm_mb.loss_type": "bce", | |
| "wm_mb.temperature": 0.1, | |
| # losses | |
| "losses": { # encodec loss + tf = 10 | |
| "adv": 4.0, | |
| "feat": 4.0, | |
| "l1": 0.1, | |
| "mel": 0.0, | |
| "msspec": 2.0, | |
| "sisnr": 0.0, | |
| "tf_loudnessratio": 10.0, | |
| }, | |
| "losses.wm_detection": 1.0, | |
| "losses.wm_mb": 1.0, | |
| } | |
| ) | |
| launcher.bind_(label="kbits16") | |
| lrs = [5e-5] | |
| seeds = [1, 2, 3, 4] | |
| with launcher.job_array(): | |
| for lr in lrs: | |
| for seed in seeds: | |
| launcher({ | |
| "optim.lr": lr, | |
| "seed": seed, | |
| }) | |