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. | |
""" | |
Evaluation with objective metrics for the pretrained AudioGen models. | |
This grid takes signature from the training grid and runs evaluation-only stage. | |
When running the grid for the first time, please use: | |
REGEN=1 dora grid audiogen.audiogen_pretrained_16khz_eval | |
and re-use the REGEN=1 option when the grid is changed to force regenerating it. | |
Note that you need the proper metrics external libraries setup to use all | |
the objective metrics activated in this grid. Refer to the README for more information. | |
""" | |
import os | |
from ..musicgen._explorers import GenerationEvalExplorer | |
from ...environment import AudioCraftEnvironment | |
from ... import train | |
def eval(launcher, batch_size: int = 32): | |
opts = { | |
'dset': 'audio/audiocaps_16khz', | |
'solver/audiogen/evaluation': 'objective_eval', | |
'execute_only': 'evaluate', | |
'+dataset.evaluate.batch_size': batch_size, | |
'+metrics.fad.tf.batch_size': 32, | |
} | |
# binary for FAD computation: replace this path with your own path | |
metrics_opts = { | |
'metrics.fad.tf.bin': '/data/home/jadecopet/local/usr/opt/google-research' | |
} | |
opt1 = {'generate.lm.use_sampling': True, 'generate.lm.top_k': 250, 'generate.lm.top_p': 0.} | |
opt2 = {'transformer_lm.two_step_cfg': True} | |
sub = launcher.bind(opts) | |
sub.bind_(metrics_opts) | |
# base objective metrics | |
sub(opt1, opt2) | |
def explorer(launcher): | |
partitions = AudioCraftEnvironment.get_slurm_partitions(['team', 'global']) | |
launcher.slurm_(gpus=4, partition=partitions) | |
if 'REGEN' not in os.environ: | |
folder = train.main.dora.dir / 'grids' / __name__.split('.', 2)[-1] | |
with launcher.job_array(): | |
for sig in folder.iterdir(): | |
if not sig.is_symlink(): | |
continue | |
xp = train.main.get_xp_from_sig(sig.name) | |
launcher(xp.argv) | |
return | |
audiogen_base = launcher.bind(solver="audiogen/audiogen_base_16khz") | |
audiogen_base.bind_({'autocast': False, 'fsdp.use': True}) | |
audiogen_base_medium = audiogen_base.bind({'continue_from': '//pretrained/facebook/audiogen-medium'}) | |
audiogen_base_medium.bind_({'model/lm/model_scale': 'medium'}) | |
eval(audiogen_base_medium, batch_size=128) | |