mirror of
https://github.com/hanzoai/mugen.git
synced 2026-11-03 01:48:40 +00:00
92 lines
2.7 KiB
Python
92 lines
2.7 KiB
Python
# 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
|
|
|
|
|
|
@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,
|
|
})
|